feat(engine): support declaration body recovery and forward-decl resolution

internal/engine/declaration_test.go

@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	"github.com/HarshK97/diffmantic/internal/testutil"
+	"github.com/HarshK97/diffmantic/internal/treesitter"
 )
 
 func TestGetDeclarationName(t *testing.T) {
@@ -330,3 +331,127 @@ func TestMatchDeclarationIntegration(t *testing.T) {
 		t.Error("declaration pre-match should map foo -> foo")
 	}
 }
+
+func TestBottomUpDeclarationNameAffinity(t *testing.T) {
+	// BottomUp should skip c1 because its declaration name doesn't match t1.
+	leafA1 := testutil.Leaf("id", "x")
+	leafA2 := testutil.Leaf("id", "x")
+	leafB1 := testutil.Leaf("id", "y")
+	leafB2 := testutil.Leaf("id", "y")
+
+	t1 := testutil.Node("function_declaration", "",
+		testutil.Leaf("identifier", "foo"),
+		leafA1,
+		leafB1,
+	)
+	t1.Language = "go"
+
+	c1 := testutil.Node("function_declaration", "",
+		testutil.Leaf("identifier", "bar"),
+		leafA2,
+	)
+	c1.Language = "go"
+
+	c2 := testutil.Node("function_declaration", "",
+		testutil.Leaf("identifier", "foo"),
+		leafB2,
+	)
+	c2.Language = "go"
+
+	m := NewMapping()
+	m.Add(leafA1, leafA2)
+	m.Add(leafB1, leafB2)
+
+	picked := candidate(t1, []*treesitter.ASTNode{c1, c2}, m)
+	if picked != c2 {
+		t.Errorf("candidate() selected %v, want c2 (%v)", picked, c2)
+	}
+}
+
+func TestMatchDeclarationsEquivalentTypes(t *testing.T) {
+	t.Run("matches function_declaration with variable_declaration in lua", func(t *testing.T) {
+		src := testutil.Node("chunk", "",
+			testutil.Node("function_declaration", "",
+				testutil.Leaf("identifier", "sync_once_impl"),
+				testutil.Node("block", "",
+					testutil.Leaf("return_statement", "return"),
+				),
+			),
+		)
+		src.Language = "lua"
+
+		dst := testutil.Node("chunk", "",
+			testutil.Node("variable_declaration", "",
+				testutil.Leaf("identifier", "sync_once_impl"),
+				testutil.Node("block", "",
+					testutil.Leaf("return_statement", "return"),
+				),
+			),
+		)
+		dst.Language = "lua"
+
+		m := NewMapping()
+		matchDeclarations(src, dst, m)
+
+		srcFn := src.Children[0]
+		dstFn := dst.Children[0]
+		if !m.Has(srcFn) {
+			t.Fatal("src function_declaration should map to dst variable_declaration under equivalent types")
+		}
+		if m.Src()[srcFn] != dstFn {
+			t.Errorf("expected %v to map to %v, got %v", srcFn, dstFn, m.Src()[srcFn])
+		}
+
+		srcBlock := srcFn.Children[1]
+		dstBlock := dstFn.Children[1]
+		if m.Src()[srcBlock] != dstBlock {
+			t.Errorf("expected body blocks to be matched and recovered")
+		}
+	})
+
+	t.Run("disambiguates forward declaration vs full definition by size", func(t *testing.T) {
+		srcForward := testutil.Node("variable_declaration", "",
+			testutil.Leaf("identifier", "sync_once_impl"),
+		)
+		srcForward.Language = "lua"
+
+		srcFull := testutil.Node("function_declaration", "",
+			testutil.Leaf("identifier", "sync_once_impl"),
+			testutil.Node("parameters", ""),
+			testutil.Node("block", "",
+				testutil.Leaf("statement", "a"),
+				testutil.Leaf("statement", "b"),
+				testutil.Leaf("statement", "c"),
+			),
+		)
+		srcFull.Language = "lua"
+
+		src := testutil.Node("chunk", "", srcForward, srcFull)
+		src.Language = "lua"
+
+		dstFull := testutil.Node("function_declaration", "",
+			testutil.Leaf("identifier", "sync_once_impl"),
+			testutil.Node("parameters", ""),
+			testutil.Node("block", "",
+				testutil.Leaf("statement", "a"),
+				testutil.Leaf("statement", "b"),
+			),
+		)
+		dstFull.Language = "lua"
+		dst := testutil.Node("chunk", "", dstFull)
+		dst.Language = "lua"
+
+		m := NewMapping()
+		matchDeclarations(src, dst, m)
+
+		if m.Has(srcForward) {
+			t.Errorf("forward declaration should not steal destination function")
+		}
+		if !m.Has(srcFull) {
+			t.Fatalf("full definition should map to destination function")
+		}
+		if m.Src()[srcFull] != dstFull {
+			t.Errorf("expected srcFull to map to dstFull")
+		}
+	})
+}

internal/engine/matcher.go

@@ -424,20 +424,90 @@ func matchDeclarations(t1Root, t2Root *treesitter.ASTNode, m *Mapping) {
 		}
 	}
 
-	for key, d2List := range t2Map {
+	visited := make(map[decKey]bool, len(t2Map))
+	for _, d2 := range t2Decs {
+		if m.HasDst(d2) {
+			continue
+		}
+		key := decKey{name: getDeclarationName(d2), rec: getReceiverTypeName(d2)}
+		if key.name == "" || visited[key] {
+			continue
+		}
+		visited[key] = true
+
+		d2List := t2Map[key]
 		if len(d2List) != 1 {
 			continue
 		}
 		d1List, ok := t1Map[key]
-		if !ok || len(d1List) != 1 {
+		if !ok || len(d1List) == 0 {
 			continue
 		}
-		d1 := d1List[0]
-		d2 := d2List[0]
-		if TypesMatch(d1.Type, d2.Type, rules) {
-			m.Add(d1, d2)
+
+		var matchingD1s []*treesitter.ASTNode
+		for _, d1 := range d1List {
+			if TypesMatch(d1.Type, d2.Type, rules) {
+				matchingD1s = append(matchingD1s, d1)
+			}
 		}
+
+		if len(matchingD1s) > 0 {
+			bestD1 := matchingD1s[0]
+			if len(matchingD1s) > 1 {
+				// If one is a forward decl and the other is a full definition,
+				// pick the one whose subtree size is closest to d2.
+				d2Size := d2.Size()
+				bestDiff := max(bestD1.Size()-d2Size, d2Size-bestD1.Size())
+				for _, cand := range matchingD1s[1:] {
+					if diff := max(cand.Size()-d2Size, d2Size-cand.Size()); diff < bestDiff {
+						bestDiff = diff
+						bestD1 = cand
+					}
+				}
+			}
+			m.Add(bestD1, d2)
+			matchDeclarationBodies(bestD1, d2, m, rules)
+		}
+	}
+}
+
+func matchDeclarationBodies(d1, d2 *treesitter.ASTNode, m *Mapping, rules *treesitter.Rules) {
+	if d1 == nil || d2 == nil || m == nil {
+		return
+	}
+	var b1, b2 *treesitter.ASTNode
+	for _, c1 := range d1.Children {
+		if isBlockNode(c1, rules) {
+			b1 = c1
+			break
+		}
+	}
+	for _, c2 := range d2.Children {
+		if isBlockNode(c2, rules) {
+			b2 = c2
+			break
+		}
+	}
+	if b1 != nil && b2 != nil {
+		if cur := m.Dst()[b2]; cur != nil && cur != b1 {
+			m.Remove(cur)
+		}
+		if cur := m.Src()[b1]; cur != nil && cur != b2 {
+			m.Remove(b1)
+		}
+		m.Add(b1, b2)
+		Recover(b1, b2, m)
+	}
+}
+
+func isBlockNode(n *treesitter.ASTNode, rules *treesitter.Rules) bool {
+	if n == nil {
+		return false
+	}
+	if rules != nil && len(rules.Blocks) > 0 {
+		return slices.Contains(rules.Blocks, n.Type)
 	}
+	return n.Type == "block"
 }
 
 func findDeclarations(root *treesitter.ASTNode) []*treesitter.ASTNode {

internal/treesitter/c/rules.yml

@@ -159,6 +159,9 @@ declarations:
 identifiers:
   - identifier
 
+blocks:
+  - compound_statement
+
 pairs:
 
   - field_designator

internal/treesitter/cpp/rules.yml

@@ -214,6 +214,9 @@ identifiers:
   - destructor_name
   - operator_name
 
+blocks:
+  - compound_statement
+
 pairs:
 
   - field_designator

internal/treesitter/go/rules.yml

@@ -130,6 +130,9 @@ identifiers:
   - type_identifier
   - package_identifier
 
+blocks:
+  - block
+
 unordered:
 
   - interface_type

internal/treesitter/java/rules.yml

@@ -190,6 +190,10 @@ declarations:
 identifiers:
   - identifier
 
+blocks:
+  - block
+  - constructor_body
+
 unordered:
 
   - annotation

internal/treesitter/javascript/rules.yml

@@ -177,6 +177,9 @@ identifiers:
   - identifier
   - private_property_identifier
 
+blocks:
+  - statement_block
+
 unordered:
 
   - object

internal/treesitter/lua/rules.yml

@@ -113,6 +113,9 @@ declarations:
 identifiers:
   - identifier
 
+blocks:
+  - block
+
 pairs:
 
   - field

internal/treesitter/php/rules.yml

@@ -175,6 +175,9 @@ identifiers:
   - name
   - variable_name
 
+blocks:
+  - compound_statement
+
 unordered:
 
   - array_creation_expression

internal/treesitter/python/rules.yml

@@ -158,6 +158,9 @@ declarations:
 identifiers:
   - identifier
 
+blocks:
+  - block
+
 unordered:
 
   - dictionary

internal/treesitter/ruby/rules.yml

@@ -159,6 +159,9 @@ identifiers:
   - identifier
   - constant
 
+blocks:
+  - block
+
 unordered:
 
   - hash

internal/treesitter/rules.go

@@ -17,6 +17,7 @@ type Rules struct {
 	Keywords        []string          `yaml:"keywords"`
 	Declarations    []string          `yaml:"declarations"`
 	Identifiers     []string          `yaml:"identifiers"`
+	Blocks          []string          `yaml:"blocks"`
 	Pairs           []string          `yaml:"pairs"`
 	Unordered       []string          `yaml:"unordered"`
 	EquivalentTypes [][]string        `yaml:"equivalent_types"`

internal/treesitter/rules_test.go

@@ -309,3 +309,36 @@ func TestEveryLanguageIdentifiersAreValidSymbols(t *testing.T) {
 		}
 	}
 }
+
+func TestEveryLanguageBlocksAreValidSymbols(t *testing.T) {
+	for _, ext := range []string{
+		"c.c", "cpp.cc", "css.css", "go.go", "html.html", "java.java",
+		"javascript.js", "json.json", "lua.lua", "php.php", "python.py",
+		"ruby.rb", "rust.rs", "toml.toml", "tsx.tsx", "typescript.ts",
+		"yaml.yaml", "zig.zig",
+	} {
+		entry := DetectGrammarEntry(ext)
+		if entry == nil {
+			continue
+		}
+		lang := entry.Language()
+		namedSymbols := make(map[string]bool)
+		for i := 0; i < int(lang.SymbolCount) && i < len(lang.SymbolNames); i++ {
+			name := lang.SymbolNames[i]
+			isNamed := i < len(lang.SymbolMetadata) && lang.SymbolMetadata[i].Named
+			if name != "" && isNamed {
+				namedSymbols[name] = true
+			}
+		}
+
+		rules := GetRules(entry.Name)
+		if rules == nil {
+			continue
+		}
+		for _, sym := range rules.Blocks {
+			if !namedSymbols[sym] {
+				t.Errorf("language %s: blocks symbol %q is not a valid named symbol in grammar", entry.Name, sym)
+			}
+		}
+	}
+}

internal/treesitter/rust/rules.yml

@@ -154,6 +154,9 @@ declarations:
 identifiers:
   - identifier
 
+blocks:
+  - block
+
 unordered:
 
   - use_declaration

internal/treesitter/tsx/rules.yml

@@ -195,6 +195,9 @@ identifiers:
   - nested_type_identifier
   - private_property_identifier
 
+blocks:
+  - statement_block
+
 unordered:
 
   - object

internal/treesitter/typescript/rules.yml

@@ -181,6 +181,9 @@ identifiers:
   - nested_type_identifier
   - private_property_identifier
 
+blocks:
+  - statement_block
+
 unordered:
 
   - object