break up GC logic

License: MIT
Signed-off-by: Jeromy <jeromyj@gmail.com>

pin/gc/gc.go

@@ -29,38 +29,9 @@ func GC(ctx context.Context, bs bstore.GCBlockstore, pn pin.Pinner) (<-chan key.
 	bsrv := bserv.New(bs, offline.Exchange(bs))
 	ds := dag.NewDAGService(bsrv)
 
-	// KeySet currently implemented in memory, in the future, may be bloom filter or
-	// disk backed to conserve memory.
-	gcs := key.NewKeySet()
-	for _, k := range pn.RecursiveKeys() {
-		gcs.Add(k)
-		nd, err := ds.Get(ctx, k)
-		if err != nil {
-			return nil, err
-		}
-
-		// EnumerateChildren recursively walks the dag and adds the keys to the given set
-		err = dag.EnumerateChildren(ctx, ds, nd, gcs)
-		if err != nil {
-			return nil, err
-		}
-	}
-	for _, k := range pn.DirectKeys() {
-		gcs.Add(k)
-	}
-	for _, k := range pn.InternalPins() {
-		gcs.Add(k)
-
-		nd, err := ds.Get(ctx, k)
-		if err != nil {
-			return nil, err
-		}
-
-		// EnumerateChildren recursively walks the dag and adds the keys to the given set
-		err = dag.EnumerateChildren(ctx, ds, nd, gcs)
-		if err != nil {
-			return nil, err
-		}
+	gcs, err := ColoredSet(pn, ds)
+	if err != nil {
+		return nil, err
 	}
 
 	keychan, err := bs.AllKeysChan(ctx)
@@ -97,3 +68,42 @@ func GC(ctx context.Context, bs bstore.GCBlockstore, pn pin.Pinner) (<-chan key.
 
 	return output, nil
 }
+
+func Descendants(ds dag.DAGService, set key.KeySet, roots []key.Key) error {
+	for _, k := range roots {
+		set.Add(k)
+		nd, err := ds.Get(context.Background(), k)
+		if err != nil {
+			return err
+		}
+
+		// EnumerateChildren recursively walks the dag and adds the keys to the given set
+		err = dag.EnumerateChildren(context.Background(), ds, nd, set)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func ColoredSet(pn pin.Pinner, ds dag.DAGService) (key.KeySet, error) {
+	// KeySet currently implemented in memory, in the future, may be bloom filter or
+	// disk backed to conserve memory.
+	gcs := key.NewKeySet()
+	err := Descendants(ds, gcs, pn.RecursiveKeys())
+	if err != nil {
+		return nil, err
+	}
+
+	for _, k := range pn.DirectKeys() {
+		gcs.Add(k)
+	}
+
+	err = Color(ds, gcs, pn.InternalPins())
+	if err != nil {
+		return nil, err
+	}
+
+	return gcs, nil
+}