7 files changed, 23 insertions, 140 deletions
diff --git a/entities/bug/bug.go b/entities/bug/bug.go
index b0f46c0b..deb00c7c 100644
--- a/entities/bug/bug.go
+++ b/entities/bug/bug.go
@@ -20,9 +20,12 @@ var _ entity.Interface = &Bug{}
 // 4: with DAG entity framework
 const formatVersion = 4
 
+const Typename = "bug"
+const Namespace = "bugs"
+
 var def = dag.Definition{
-	Typename:             "bug",
-	Namespace:            "bugs",
+	Typename:             Typename,
+	Namespace:            Namespace,
 	OperationUnmarshaler: operationUnmarshaler,
 	FormatVersion:        formatVersion,
 }
@@ -42,9 +45,11 @@ type Bug struct {
 
 // NewBug create a new Bug
 func NewBug() *Bug {
-	return &Bug{
-		Entity: dag.New(def),
-	}
+	return wrapper(dag.New(def))
+}
+
+func wrapper(e *dag.Entity) *Bug {
+	return &Bug{Entity: e}
 }
 
 func simpleResolvers(repo repository.ClockedRepo) entity.Resolvers {
@@ -60,49 +65,17 @@ func Read(repo repository.ClockedRepo, id entity.Id) (*Bug, error) {
 
 // ReadWithResolver will read a bug from its Id, with custom resolvers
 func ReadWithResolver(repo repository.ClockedRepo, resolvers entity.Resolvers, id entity.Id) (*Bug, error) {
-	e, err := dag.Read(def, repo, resolvers, id)
-	if err != nil {
-		return nil, err
-	}
-	return &Bug{Entity: e}, nil
-}
-
-type StreamedBug struct {
-	Bug *Bug
-	Err error
+	return dag.Read(def, wrapper, repo, resolvers, id)
 }
 
 // ReadAll read and parse all local bugs
-func ReadAll(repo repository.ClockedRepo) <-chan StreamedBug {
-	return readAll(repo, simpleResolvers(repo))
+func ReadAll(repo repository.ClockedRepo) <-chan entity.StreamedEntity[*Bug] {
+	return dag.ReadAll(def, wrapper, repo, simpleResolvers(repo))
 }
 
 // ReadAllWithResolver read and parse all local bugs
-func ReadAllWithResolver(repo repository.ClockedRepo, resolvers entity.Resolvers) <-chan StreamedBug {
-	return readAll(repo, resolvers)
-}
-
-// Read and parse all available bug with a given ref prefix
-func readAll(repo repository.ClockedRepo, resolvers entity.Resolvers) <-chan StreamedBug {
-	out := make(chan StreamedBug)
-
-	go func() {
-		defer close(out)
-
-		for streamedEntity := range dag.ReadAll(def, repo, resolvers) {
-			if streamedEntity.Err != nil {
-				out <- StreamedBug{
-					Err: streamedEntity.Err,
-				}
-			} else {
-				out <- StreamedBug{
-					Bug: &Bug{Entity: streamedEntity.Entity},
-				}
-			}
-		}
-	}()
-
-	return out
+func ReadAllWithResolver(repo repository.ClockedRepo, resolvers entity.Resolvers) <-chan entity.StreamedEntity[*Bug] {
+	return dag.ReadAll(def, wrapper, repo, resolvers)
 }
 
 // ListLocalIds list all the available local bug ids
diff --git a/entities/bug/bug_actions.go b/entities/bug/bug_actions.go
index c25b9243..198e4ed0 100644
--- a/entities/bug/bug_actions.go
+++ b/entities/bug/bug_actions.go
@@ -23,33 +23,14 @@ func Push(repo repository.Repo, remote string) (string, error) {
 // Note: an author is necessary for the case where a merge commit is created, as this commit will
 // have an author and may be signed if a signing key is available.
 func Pull(repo repository.ClockedRepo, resolvers entity.Resolvers, remote string, mergeAuthor identity.Interface) error {
-	return dag.Pull(def, repo, resolvers, remote, mergeAuthor)
+	return dag.Pull(def, wrapper, repo, resolvers, remote, mergeAuthor)
 }
 
 // MergeAll will merge all the available remote bug
 // Note: an author is necessary for the case where a merge commit is created, as this commit will
 // have an author and may be signed if a signing key is available.
 func MergeAll(repo repository.ClockedRepo, resolvers entity.Resolvers, remote string, mergeAuthor identity.Interface) <-chan entity.MergeResult {
-	out := make(chan entity.MergeResult)
-
-	go func() {
-		defer close(out)
-
-		results := dag.MergeAll(def, repo, resolvers, remote, mergeAuthor)
-
-		// wrap the dag.Entity into a complete Bug
-		for result := range results {
-			result := result
-			if result.Entity != nil {
-				result.Entity = &Bug{
-					Entity: result.Entity.(*dag.Entity),
-				}
-			}
-			out <- result
-		}
-	}()
-
-	return out
+	return dag.MergeAll(def, wrapper, repo, resolvers, remote, mergeAuthor)
 }
 
 // Remove will remove a local bug from its entity.Id
diff --git a/entities/bug/err.go b/entities/bug/err.go
deleted file mode 100644
index 1bd174bb..00000000
--- a/entities/bug/err.go
+++ /dev/null
@@ -1,17 +0,0 @@
-package bug
-
-import (
-	"errors"
-
-	"github.com/MichaelMure/git-bug/entity"
-)
-
-var ErrBugNotExist = errors.New("bug doesn't exist")
-
-func NewErrMultipleMatchBug(matching []entity.Id) *entity.ErrMultipleMatch {
-	return entity.NewErrMultipleMatch("bug", matching)
-}
-
-func NewErrMultipleMatchOp(matching []entity.Id) *entity.ErrMultipleMatch {
-	return entity.NewErrMultipleMatch("operation", matching)
-}
diff --git a/entities/bug/operation.go b/entities/bug/operation.go
index 41d80700..04365046 100644
--- a/entities/bug/operation.go
+++ b/entities/bug/operation.go
@@ -21,12 +21,7 @@ const (
 )
 
 // Operation define the interface to fulfill for an edit operation of a Bug
-type Operation interface {
-	dag.Operation
-
-	// Apply the operation to a Snapshot to create the final state
-	Apply(snapshot *Snapshot)
-}
+type Operation = dag.OperationWithApply[*Snapshot]
 
 // make sure that package external operations do conform to our interface
 var _ Operation = &dag.NoOpOperation[*Snapshot]{}
diff --git a/entities/bug/resolver.go b/entities/bug/resolver.go
index e7beb0e4..b0a05917 100644
--- a/entities/bug/resolver.go
+++ b/entities/bug/resolver.go
@@ -16,6 +16,6 @@ func NewSimpleResolver(repo repository.ClockedRepo) *SimpleResolver {
 	return &SimpleResolver{repo: repo}
 }
 
-func (r *SimpleResolver) Resolve(id entity.Id) (entity.Interface, error) {
+func (r *SimpleResolver) Resolve(id entity.Id) (entity.Resolved, error) {
 	return Read(r.repo, id)
 }
diff --git a/entities/bug/snapshot.go b/entities/bug/snapshot.go
index 333fe207..5c260d85 100644
--- a/entities/bug/snapshot.go
+++ b/entities/bug/snapshot.go
@@ -43,6 +43,10 @@ func (snap *Snapshot) AllOperations() []dag.Operation {
 	return snap.Operations
 }
 
+func (snap *Snapshot) AppendOperation(op dag.Operation) {
+	snap.Operations = append(snap.Operations, op)
+}
+
 // EditTime returns the last time a bug was modified
 func (snap *Snapshot) EditTime() time.Time {
 	if len(snap.Operations) == 0 {
diff --git a/entities/bug/with_snapshot.go b/entities/bug/with_snapshot.go
deleted file mode 100644
index 0474cac7..00000000
--- a/entities/bug/with_snapshot.go
+++ /dev/null
@@ -1,53 +0,0 @@
-package bug
-
-import (
-	"github.com/MichaelMure/git-bug/repository"
-)
-
-var _ Interface = &WithSnapshot{}
-
-// WithSnapshot encapsulate a Bug and maintain the corresponding Snapshot efficiently
-type WithSnapshot struct {
-	*Bug
-	snap *Snapshot
-}
-
-func (b *WithSnapshot) Compile() *Snapshot {
-	if b.snap == nil {
-		snap := b.Bug.Compile()
-		b.snap = snap
-	}
-	return b.snap
-}
-
-// Append intercept Bug.Append() to update the snapshot efficiently
-func (b *WithSnapshot) Append(op Operation) {
-	b.Bug.Append(op)
-
-	if b.snap == nil {
-		return
-	}
-
-	op.Apply(b.snap)
-	b.snap.Operations = append(b.snap.Operations, op)
-}
-
-// Commit intercept Bug.Commit() to update the snapshot efficiently
-func (b *WithSnapshot) Commit(repo repository.ClockedRepo) error {
-	err := b.Bug.Commit(repo)
-
-	if err != nil {
-		b.snap = nil
-		return err
-	}
-
-	// Commit() shouldn't change anything of the bug state apart from the
-	// initial ID set
-
-	if b.snap == nil {
-		return nil
-	}
-
-	b.snap.id = b.Bug.Id()
-	return nil
-}