Avoid repeated prefix scans when projecting item policies

internal/engine/imports.go

@@ -961,17 +961,28 @@ func projectNestedCollectionPolicies(composed *composedDocument, scope acquisiti
 		return
 	}
 	baseOffsets := make(map[string]int)
+	var prefixLengths map[string]int
 	for sourceIndex, contribution := range scope.contributions {
+		if prefixLengths != nil && sourceIndex > 0 {
+			addSequencePrefixLengths(scope.contributions[sourceIndex-1].document, Path{}, prefixLengths)
+		}
 		if contribution.nested == nil || len(contribution.nested.collectionPolicies.collections) == 0 {
 			continue
 		}
-		itemIndex := priorSequenceLength([]acquisitionScope{{contributions: scope.contributions[:sourceIndex]}}, contribution.destination)
-		if assigned, found := ranges[sequenceRangeKey{sourceIndex: sourceIndex, path: contribution.destination.String()}]; found {
-			if projected, ok := assigned.outputForInput(0); ok {
-				itemIndex = projected
-			}
-		}
 		key := contribution.destination.String()
+		assigned := ranges[sequenceRangeKey{sourceIndex: sourceIndex, path: key}]
+		itemIndex, found := assigned.outputForInput(0)
+		if !found {
+			// Normal composition supplies the assignment directly. For a
+			// handoff without ranges, scan each preceding tree only once.
+			if prefixLengths == nil {
+				prefixLengths = make(map[string]int)
+				for _, previous := range scope.contributions[:sourceIndex] {
+					addSequencePrefixLengths(previous.document, Path{}, prefixLengths)
+				}
+			}
+			itemIndex = prefixLengths[key]
+		}
 		base, initialized := baseOffsets[key]
 		policy, _ := composed.collectionPolicies.at(contribution.destination)
 		if !initialized {
@@ -986,6 +997,20 @@ func projectNestedCollectionPolicies(composed *composedDocument, scope acquisiti
 	}
 }
 
+func addSequencePrefixLengths(node *document.Node, path Path, lengths map[string]int) {
+	if node == nil {
+		return
+	}
+	if node.Kind == document.Sequence {
+		lengths[path.String()] += len(node.Sequence)
+		return
+	}
+	// Contribution destinations select mapping paths, not sequence members.
+	for _, entry := range node.Mapping {
+		addSequencePrefixLengths(entry.Value, path.Key(entry.Key), lengths)
+	}
+}
+
 func projectNestedOrigins(target map[string]Origin, root *document.Node, origins map[string]Origin, prefix Path) {
 	var visit func(*document.Node, Path)
 	visit = func(node *document.Node, path Path) {

internal/engine/item_policy_test.go

@@ -195,6 +195,40 @@ func TestNestedItemPolicyProjectionUsesStableOutputIndexes(t *testing.T) {
 	}
 }
 
+func TestItemPolicyProjectionCountsContributionsWithoutLocalPolicy(t *testing.T) {
+	t.Parallel()
+	items := Keys("items")
+	for _, assignments := range []map[int]int{nil, {1: 2, 3: 4}, {1: 2}, {3: 4}} {
+		t.Run(fmt.Sprint(assignments), func(t *testing.T) {
+			local := collectionPolicyScope{}
+			local.set(Keys("steps"), collectionPolicy{identityFields: []string{"name"}, identityPresent: true})
+			scope := acquisitionScope{contributions: []acquiredContribution{
+				{document: parseItemPolicyDocument(t, "inline", "items: [{id: inline-a}, {id: inline-b}]\n")},
+				{document: parseItemPolicyDocument(t, "first", "items: [{id: first}]\n"), destination: items, nested: &composedDocument{collectionPolicies: local}},
+				{document: parseItemPolicyDocument(t, "gap", "items: [{id: gap}]\n"), destination: items},
+				{document: parseItemPolicyDocument(t, "second", "items: [{id: second}]\n"), destination: items, nested: &composedDocument{collectionPolicies: local}},
+			}}
+			ranges := make(map[sequenceRangeKey]sequenceRange)
+			for sourceIndex, output := range assignments {
+				ranges[sequenceRangeKey{sourceIndex: sourceIndex, path: items.String()}] = sequenceRange{start: output, end: output + 1}
+			}
+			composed := &composedDocument{}
+			projectNestedCollectionPolicies(composed, scope, ranges)
+			policy, _ := composed.collectionPolicies.at(items)
+			if len(policy.recordOrder) != 5 {
+				t.Fatalf("record policies = %d, want 5 positions", len(policy.recordOrder))
+			}
+			for index := range 5 {
+				record, _ := policy.recordPolicyAt(index)
+				_, found := record.scope.at(Keys("steps"))
+				if found != (index == 2 || index == 4) {
+					t.Errorf("local policy at item %d = %t", index, found)
+				}
+			}
+		})
+	}
+}
+
 func TestOverlaySequenceTransfersIncomingItemPolicyOwnership(t *testing.T) {
 	path := Keys("records")
 	registered := mustCompileSequenceIdentityPolicy(t, []sequenceIdentityRegistration{{selector: sequenceIdentitySelector{literalSequenceSelector("records")}, identityFields: []string{"id"}}})