4 files changed, 261 insertions, 125 deletions

diff --git a/libpod/image/filters.go b/libpod/image/filters.go
index 9738a7d5e..db647954f 100644
--- a/libpod/image/filters.go
+++ b/libpod/image/filters.go
@@ -29,6 +29,26 @@ func CreatedBeforeFilter(createTime time.Time) ResultFilter {
 	}
 }
 
+// IntermediateFilter returns filter for intermediate images (i.e., images
+// with children and no tags).
+func (ir *Runtime) IntermediateFilter(ctx context.Context, images []*Image) (ResultFilter, error) {
+	tree, err := ir.layerTree()
+	if err != nil {
+		return nil, err
+	}
+	return func(i *Image) bool {
+		if len(i.Names()) > 0 {
+			return true
+		}
+		children, err := tree.children(ctx, i, false)
+		if err != nil {
+			logrus.Error(err.Error())
+			return false
+		}
+		return len(children) == 0
+	}, nil
+}
+
 // CreatedAfterFilter allows you to filter on images created after
 // the given time.Time
 func CreatedAfterFilter(createTime time.Time) ResultFilter {
diff --git a/libpod/image/image.go b/libpod/image/image.go
index 4664da63e..6106084d5 100644
--- a/libpod/image/image.go
+++ b/libpod/image/image.go
@@ -859,26 +859,6 @@ func (i *Image) Dangling() bool {
 	return len(i.Names()) == 0
 }
 
-// Intermediate returns true if the image is cache or intermediate image.
-// Cache image has parent and child.
-func (i *Image) Intermediate(ctx context.Context) (bool, error) {
-	parent, err := i.IsParent(ctx)
-	if err != nil {
-		return false, err
-	}
-	if !parent {
-		return false, nil
-	}
-	img, err := i.GetParent(ctx)
-	if err != nil {
-		return false, err
-	}
-	if img != nil {
-		return true, nil
-	}
-	return false, nil
-}
-
 // User returns the image's user
 func (i *Image) User(ctx context.Context) (string, error) {
 	imgInspect, err := i.inspect(ctx, false)
@@ -1217,7 +1197,7 @@ func splitString(input string) string {
 // the parent of any other layer in store. Double check that image with that
 // layer exists as well.
 func (i *Image) IsParent(ctx context.Context) (bool, error) {
-	children, err := i.getChildren(ctx, 1)
+	children, err := i.getChildren(ctx, false)
 	if err != nil {
 		if errors.Cause(err) == ErrImageIsBareList {
 			return false, nil
@@ -1292,63 +1272,16 @@ func areParentAndChild(parent, child *imgspecv1.Image) bool {
 
 // GetParent returns the image ID of the parent. Return nil if a parent is not found.
 func (i *Image) GetParent(ctx context.Context) (*Image, error) {
-	var childLayer *storage.Layer
-	images, err := i.imageruntime.GetImages()
+	tree, err := i.imageruntime.layerTree()
 	if err != nil {
 		return nil, err
 	}
-	if i.TopLayer() != "" {
-		if childLayer, err = i.imageruntime.store.Layer(i.TopLayer()); err != nil {
-			return nil, err
-		}
-	}
-	// fetch the configuration for the child image
-	child, err := i.ociv1Image(ctx)
-	if err != nil {
-		if errors.Cause(err) == ErrImageIsBareList {
-			return nil, nil
-		}
-		return nil, err
-	}
-	for _, img := range images {
-		if img.ID() == i.ID() {
-			continue
-		}
-		candidateLayer := img.TopLayer()
-		// as a child, our top layer, if we have one, is either the
-		// candidate parent's layer, or one that's derived from it, so
-		// skip over any candidate image where we know that isn't the
-		// case
-		if childLayer != nil {
-			// The child has at least one layer, so a parent would
-			// have a top layer that's either the same as the child's
-			// top layer or the top layer's recorded parent layer,
-			// which could be an empty value.
-			if candidateLayer != childLayer.Parent && candidateLayer != childLayer.ID {
-				continue
-			}
-		} else {
-			// The child has no layers, but the candidate does.
-			if candidateLayer != "" {
-				continue
-			}
-		}
-		// fetch the configuration for the candidate image
-		candidate, err := img.ociv1Image(ctx)
-		if err != nil {
-			return nil, err
-		}
-		// compare them
-		if areParentAndChild(candidate, child) {
-			return img, nil
-		}
-	}
-	return nil, nil
+	return tree.parent(ctx, i)
 }
 
 // GetChildren returns a list of the imageIDs that depend on the image
 func (i *Image) GetChildren(ctx context.Context) ([]string, error) {
-	children, err := i.getChildren(ctx, 0)
+	children, err := i.getChildren(ctx, true)
 	if err != nil {
 		if errors.Cause(err) == ErrImageIsBareList {
 			return nil, nil
@@ -1358,62 +1291,15 @@ func (i *Image) GetChildren(ctx context.Context) ([]string, error) {
 	return children, nil
 }
 
-// getChildren returns a list of at most "max" imageIDs that depend on the image
-func (i *Image) getChildren(ctx context.Context, max int) ([]string, error) {
-	var children []string
-
-	if _, err := i.toImageRef(ctx); err != nil {
-		return nil, nil
-	}
-
-	images, err := i.imageruntime.GetImages()
-	if err != nil {
-		return nil, err
-	}
-
-	// fetch the configuration for the parent image
-	parent, err := i.ociv1Image(ctx)
+// getChildren returns a list of imageIDs that depend on the image. If all is
+// false, only the first child image is returned.
+func (i *Image) getChildren(ctx context.Context, all bool) ([]string, error) {
+	tree, err := i.imageruntime.layerTree()
 	if err != nil {
 		return nil, err
 	}
-	parentLayer := i.TopLayer()
 
-	for _, img := range images {
-		if img.ID() == i.ID() {
-			continue
-		}
-		if img.TopLayer() == "" {
-			if parentLayer != "" {
-				// this image has no layers, but we do, so
-				// it can't be derived from this one
-				continue
-			}
-		} else {
-			candidateLayer, err := img.Layer()
-			if err != nil {
-				return nil, err
-			}
-			// if this image's top layer is not our top layer, and is not
-			// based on our top layer, we can skip it
-			if candidateLayer.Parent != parentLayer && candidateLayer.ID != parentLayer {
-				continue
-			}
-		}
-		// fetch the configuration for the candidate image
-		candidate, err := img.ociv1Image(ctx)
-		if err != nil {
-			return nil, err
-		}
-		// compare them
-		if areParentAndChild(parent, candidate) {
-			children = append(children, img.ID())
-		}
-		// if we're not building an exhaustive list, maybe we're done?
-		if max > 0 && len(children) >= max {
-			break
-		}
-	}
-	return children, nil
+	return tree.children(ctx, i, all)
 }
 
 // InputIsID returns a bool if the user input for an image
@@ -1670,6 +1556,7 @@ type LayerInfo struct {
 
 // GetLayersMapWithImageInfo returns map of image-layers, with associated information like RepoTags, parent and list of child layers.
 func GetLayersMapWithImageInfo(imageruntime *Runtime) (map[string]*LayerInfo, error) {
+	// TODO: evaluate if we can reuse `layerTree` here.
 
 	// Memory allocated to store map of layers with key LayerID.
 	// Map will build dependency chain with ParentID and ChildID(s)
diff --git a/libpod/image/layer_tree.go b/libpod/image/layer_tree.go
new file mode 100644
index 000000000..3699655fd
--- /dev/null
+++ b/libpod/image/layer_tree.go
@@ -0,0 +1,222 @@
+package image
+
+import (
+	"context"
+
+	ociv1 "github.com/opencontainers/image-spec/specs-go/v1"
+	"github.com/pkg/errors"
+)
+
+// layerTree is an internal representation of local layers.
+type layerTree struct {
+	// nodes is the actual layer tree with layer IDs being keys.
+	nodes map[string]*layerNode
+	// ociCache is a cache for Image.ID -> OCI Image. Translations are done
+	// on-demand.
+	ociCache map[string]*ociv1.Image
+}
+
+// node returns a layerNode for the specified layerID.
+func (t *layerTree) node(layerID string) *layerNode {
+	node, exists := t.nodes[layerID]
+	if !exists {
+		node = &layerNode{}
+		t.nodes[layerID] = node
+	}
+	return node
+}
+
+// toOCI returns an OCI image for the specified image.
+func (t *layerTree) toOCI(ctx context.Context, i *Image) (*ociv1.Image, error) {
+	var err error
+	oci, exists := t.ociCache[i.ID()]
+	if !exists {
+		oci, err = i.ociv1Image(ctx)
+		t.ociCache[i.ID()] = oci
+	}
+	return oci, err
+}
+
+// layerNode is a node in a layerTree.  It's ID is the key in a layerTree.
+type layerNode struct {
+	children []*layerNode
+	images   []*Image
+	parent   *layerNode
+}
+
+// layerTree extracts a layerTree from the layers in the local storage and
+// relates them to the specified images.
+func (ir *Runtime) layerTree() (*layerTree, error) {
+	layers, err := ir.store.Layers()
+	if err != nil {
+		return nil, err
+	}
+
+	images, err := ir.GetImages()
+	if err != nil {
+		return nil, err
+	}
+
+	tree := layerTree{
+		nodes:    make(map[string]*layerNode),
+		ociCache: make(map[string]*ociv1.Image),
+	}
+
+	// First build a tree purely based on layer information.
+	for _, layer := range layers {
+		node := tree.node(layer.ID)
+		if layer.Parent == "" {
+			continue
+		}
+		parent := tree.node(layer.Parent)
+		node.parent = parent
+		parent.children = append(parent.children, node)
+	}
+
+	// Now assign the images to each (top) layer.
+	for i := range images {
+		img := images[i] // do not leak loop variable outside the scope
+		topLayer := img.TopLayer()
+		if topLayer == "" {
+			continue
+		}
+		node, exists := tree.nodes[topLayer]
+		if !exists {
+			return nil, errors.Errorf("top layer %s of image %s not found in layer tree", img.TopLayer(), img.ID())
+		}
+		node.images = append(node.images, img)
+	}
+
+	return &tree, nil
+}
+
+// children returns the image IDs of children . Child images are images
+// with either the same top layer as parent or parent being the true parent
+// layer.  Furthermore, the history of the parent and child images must match
+// with the parent having one history item less.
+// If all is true, all images are returned.  Otherwise, the first image is
+// returned.
+func (t *layerTree) children(ctx context.Context, parent *Image, all bool) ([]string, error) {
+	if parent.TopLayer() == "" {
+		return nil, nil
+	}
+
+	var children []string
+
+	parentNode, exists := t.nodes[parent.TopLayer()]
+	if !exists {
+		return nil, errors.Errorf("layer not found in layer tree: %q", parent.TopLayer())
+	}
+
+	parentID := parent.ID()
+	parentOCI, err := t.toOCI(ctx, parent)
+	if err != nil {
+		return nil, err
+	}
+
+	// checkParent returns true if child and parent are in such a relation.
+	checkParent := func(child *Image) (bool, error) {
+		if parentID == child.ID() {
+			return false, nil
+		}
+		childOCI, err := t.toOCI(ctx, child)
+		if err != nil {
+			return false, err
+		}
+		// History check.
+		return areParentAndChild(parentOCI, childOCI), nil
+	}
+
+	// addChildrenFrom adds child images of parent to children.  Returns
+	// true if any image is a child of parent.
+	addChildrenFromNode := func(node *layerNode) (bool, error) {
+		foundChildren := false
+		for _, childImage := range node.images {
+			isChild, err := checkParent(childImage)
+			if err != nil {
+				return foundChildren, err
+			}
+			if isChild {
+				foundChildren = true
+				children = append(children, childImage.ID())
+				if all {
+					return foundChildren, nil
+				}
+			}
+		}
+		return foundChildren, nil
+	}
+
+	// First check images where parent's top layer is also the parent
+	// layer.
+	for _, childNode := range parentNode.children {
+		found, err := addChildrenFromNode(childNode)
+		if err != nil {
+			return nil, err
+		}
+		if found && all {
+			return children, nil
+		}
+	}
+
+	// Now check images with the same top layer.
+	if _, err := addChildrenFromNode(parentNode); err != nil {
+		return nil, err
+	}
+
+	return children, nil
+}
+
+// parent returns the parent image or nil if no parent image could be found.
+func (t *layerTree) parent(ctx context.Context, child *Image) (*Image, error) {
+	if child.TopLayer() == "" {
+		return nil, nil
+	}
+
+	node, exists := t.nodes[child.TopLayer()]
+	if !exists {
+		return nil, errors.Errorf("layer not found in layer tree: %q", child.TopLayer())
+	}
+
+	childOCI, err := t.toOCI(ctx, child)
+	if err != nil {
+		return nil, err
+	}
+
+	// Check images from the parent node (i.e., parent layer) and images
+	// with the same layer (i.e., same top layer).
+	childID := child.ID()
+	images := node.images
+	if node.parent != nil {
+		images = append(images, node.parent.images...)
+	}
+	for _, parent := range images {
+		if parent.ID() == childID {
+			continue
+		}
+		parentOCI, err := t.toOCI(ctx, parent)
+		if err != nil {
+			return nil, err
+		}
+		// History check.
+		if areParentAndChild(parentOCI, childOCI) {
+			return parent, nil
+		}
+	}
+
+	return nil, nil
+}
+
+// hasChildrenAndParent returns true if the specified image has children and a
+// parent.
+func (t *layerTree) hasChildrenAndParent(ctx context.Context, i *Image) (bool, error) {
+	children, err := t.children(ctx, i, false)
+	if err != nil {
+		return false, err
+	}
+	if len(children) == 0 {
+		return false, nil
+	}
+	parent, err := t.parent(ctx, i)
+	return parent != nil, err
+}
diff --git a/libpod/image/prune.go b/libpod/image/prune.go
index 8c9267650..5a9ca5d8e 100644
--- a/libpod/image/prune.go
+++ b/libpod/image/prune.go
@@ -66,6 +66,12 @@ func (ir *Runtime) GetPruneImages(ctx context.Context, all bool, filterFuncs []I
 	if err != nil {
 		return nil, err
 	}
+
+	tree, err := ir.layerTree()
+	if err != nil {
+		return nil, err
+	}
+
 	for _, i := range allImages {
 		// filter the images based on this.
 		for _, filterFunc := range filterFuncs {
@@ -85,8 +91,9 @@ func (ir *Runtime) GetPruneImages(ctx context.Context, all bool, filterFuncs []I
 			}
 		}
 
-		//skip the cache or intermediate images
-		intermediate, err := i.Intermediate(ctx)
+		// skip the cache (i.e., with parent) and intermediate (i.e.,
+		// with children) images
+		intermediate, err := tree.hasChildrenAndParent(ctx, i)
 		if err != nil {
 			return nil, err
 		}