traversal: implement monotonically decrementing budgets.

It's optional, and currently not on by default, but easy to set.

traversal/fns.go

@@ -36,6 +36,7 @@ type Progress struct {
 		Path datamodel.Path
 		Link datamodel.Link
 	}
+	Budget *Budget // If present, tracks "budgets" for how many more steps we're willing to take before we should halt.
 }
 
 type Config struct {
@@ -44,6 +45,18 @@ type Config struct {
 	LinkTargetNodePrototypeChooser LinkTargetNodePrototypeChooser // Chooser for Node implementations to produce during automatic link traversal.
 }
 
+type Budget struct {
+	// Fields below are described as "monotonically-decrementing", because that's what the traversal library will do with them,
+	// but they are user-accessable and can be reset to higher numbers again by code in the visitor callbacks.  This is not recommended (why?), but possible.
+
+	// If you set any budgets (by having a non-nil Progress.Budget field), you must set some value for all of them.
+	// Traversal halts when _any_ of the budgets reaches zero.
+	// The max value of an int (math.MaxInt64) is acceptable for any budget you don't care about.
+
+	NodeBudget int64 // A monotonically-decrementing "budget" for how many more nodes we're willing to visit before halting.
+	LinkBudget int64 // A monotonically-decrementing "budget" for how many more links we're willing to load before halting.  (This is not aware of any caching; it's purely in terms of links encountered and traversed.)
+}
+
 // LinkTargetNodePrototypeChooser is a function that returns a NodePrototype based on
 // the information in a Link and/or its LinkContext.
 //

traversal/focus.go

@@ -87,6 +87,13 @@ func (prog *Progress) get(n datamodel.Node, p datamodel.Path, trackProgress bool
 	segments := p.Segments()
 	var prev datamodel.Node // for LinkContext
 	for i, seg := range segments {
+		// Check the budget!
+		if prog.Budget != nil {
+			prog.Budget.NodeBudget--
+			if prog.Budget.NodeBudget <= 0 {
+				return nil, fmt.Errorf("traversal budget for nodes visited exceeded")
+			}
+		}
 		// Traverse the segment.
 		switch n.Kind() {
 		case datamodel.Kind_Invalid:
@@ -112,6 +119,14 @@ func (prog *Progress) get(n datamodel.Node, p datamodel.Path, trackProgress bool
 		}
 		// Dereference any links.
 		for n.Kind() == datamodel.Kind_Link {
+			// Check the budget!
+			if prog.Budget != nil {
+				prog.Budget.LinkBudget--
+				if prog.Budget.LinkBudget <= 0 {
+					return nil, fmt.Errorf("traversal budget for links exceeded")
+				}
+			}
+			// Put together the context info we'll offer to the loader and prototypeChooser.
 			lnk, _ := n.AsLink()
 			lnkCtx := linking.LinkContext{
 				Ctx:        prog.Cfg.Ctx,
@@ -201,6 +216,13 @@ func (prog Progress) focusedTransform(n datamodel.Node, na datamodel.NodeAssembl
 		return na.AssignNode(n2)
 	}
 	seg, p2 := p.Shift()
+	// Check the budget!
+	if prog.Budget != nil {
+		prog.Budget.NodeBudget--
+		if prog.Budget.NodeBudget <= 0 {
+			return fmt.Errorf("traversal budget for nodes visited exceeded")
+		}
+	}
 	// Special branch for if we've entered createParent mode in an earlier step.
 	//  This needs slightly different logic because there's no prior node to reference
 	//   (and we wouldn't want to waste time creating a dummy one).
@@ -319,6 +341,14 @@ func (prog Progress) focusedTransform(n datamodel.Node, na datamodel.NodeAssembl
 		}
 		return la.Finish()
 	case datamodel.Kind_Link:
+		// Check the budget!
+		if prog.Budget != nil {
+			prog.Budget.LinkBudget--
+			if prog.Budget.LinkBudget <= 0 {
+				return fmt.Errorf("traversal budget for links exceeded")
+			}
+		}
+		// Put together the context info we'll offer to the loader and prototypeChooser.
 		lnkCtx := linking.LinkContext{
 			Ctx:        prog.Cfg.Ctx,
 			LinkPath:   prog.Path,

traversal/walk.go

@@ -86,6 +86,14 @@ func (prog Progress) WalkAdv(n datamodel.Node, s selector.Selector, fn AdvVisitF
 }
 
 func (prog Progress) walkAdv(n datamodel.Node, s selector.Selector, fn AdvVisitFn) error {
+	// Check the budget!
+	if prog.Budget != nil {
+		prog.Budget.NodeBudget--
+		if prog.Budget.NodeBudget <= 0 {
+			return fmt.Errorf("traversal budget for nodes visited exceeded")
+		}
+	}
+	// Decide if this node is matched -- do callbacks as appropriate.
 	if s.Decide(n) {
 		if err := fn(prog, n, VisitReason_SelectionMatch); err != nil {
 			return err
@@ -95,12 +103,14 @@ func (prog Progress) walkAdv(n datamodel.Node, s selector.Selector, fn AdvVisitF
 			return err
 		}
 	}
+	// If we're handling scalars (e.g. not maps and lists) we can return now.
 	nk := n.Kind()
 	switch nk {
 	case datamodel.Kind_Map, datamodel.Kind_List: // continue
 	default:
 		return nil
 	}
+	// For maps and lists: recurse (in one of two ways, depending on if the selector also states specific interests).
 	attn := s.Interests()
 	if attn == nil {
 		return prog.walkAdv_iterateAll(n, s, fn)
@@ -180,6 +190,14 @@ func (prog Progress) walkAdv_iterateSelective(n datamodel.Node, attn []datamodel
 }
 
 func (prog Progress) loadLink(v datamodel.Node, parent datamodel.Node) (datamodel.Node, error) {
+	// Check the budget!
+	if prog.Budget != nil {
+		prog.Budget.LinkBudget--
+		if prog.Budget.LinkBudget <= 0 {
+			return nil, fmt.Errorf("traversal budget for links exceeded")
+		}
+	}
+	// Put together the context info we'll offer to the loader and prototypeChooser.
 	lnk, err := v.AsLink()
 	if err != nil {
 		return nil, err