perf(treesitter): compile AST rules into lookup sets and cache pointers

CHANGELOG.md

@@ -23,6 +23,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed false-positive matches between deleted and surviving variables when expressions have the exact same shape.
 
 ### Performance
+- Compiled Tree-sitter AST transformation rules into O(1) hash sets and cached rule pointers, eliminating repeated slice scans and struct copies on every AST node (PR #118).
 - Sped up declaration matching across files with many functions or components, cutting matching time by up to 20% on large refactors (PR #108).
 - Made Git diff computations run asynchronously during status refresh in the TUI, preventing UI freezes while background diffs compute (PR #106).
 

internal/engine/matcher.go

@@ -16,6 +16,12 @@ type MatchResult struct {
 func Match(t1, t2 *treesitter.ASTNode, srcA, srcB []byte, part *LinePartition) *MatchResult {
 	mappings := NewMapping()
 
+	if t1 != nil && t2 != nil && Isomorphic(t1, t2) {
+		addIsomorphicPairs(t1, t2, mappings)
+		sortMappingsByPreOrder(mappings)
+		return &MatchResult{Mappings: mappings}
+	}
+
 	if part == nil {
 		part = NewLinePartition(srcA, srcB)
 	}
@@ -505,7 +511,7 @@ func isBlockNode(n *treesitter.ASTNode, rules *treesitter.Rules) bool {
 		return false
 	}
 	if rules != nil && len(rules.Blocks) > 0 {
-		return slices.Contains(rules.Blocks, n.Type)
+		return rules.IsBlock(n.Type)
 	}
 	return n.Type == "block"
 }

internal/treesitter/ast.go

@@ -154,7 +154,7 @@ func buildASTWithRules(n *gotreesitter.Node, src []byte, lang *gotreesitter.Lang
 		label = strings.TrimSpace(string(src[start:end]))
 	}
 
-	if isIgnored(nodeType, label, rules.Ignored) {
+	if rules.IsIgnored(nodeType, label) {
 		return nil
 	}
 
@@ -174,22 +174,17 @@ func buildASTWithRules(n *gotreesitter.Node, src []byte, lang *gotreesitter.Lang
 		node.Label = label
 	}
 
-	if rules != nil {
-		if alias, ok := rules.Aliased[nodeType]; ok {
-			node.Type = alias
-		}
-		if alias, ok := rules.Aliased[label]; ok {
-			node.Type = alias
-		}
-		if slices.Contains(rules.LabelIgnored, node.Type) {
-			node.Label = ""
-		}
-		if slices.Contains(rules.Keywords, nodeType) || (label != "" && slices.Contains(rules.Keywords, label)) {
-			node.IsKeyword = true
-		}
-		if slices.Contains(rules.Unordered, node.Type) {
-			node.IsUnordered = true
-		}
+	if alias, ok := rules.Alias(nodeType, label); ok {
+		node.Type = alias
+	}
+	if rules.IsLabelIgnored(node.Type) {
+		node.Label = ""
+	}
+	if rules.IsKeyword(nodeType, label) {
+		node.IsKeyword = true
+	}
+	if rules.IsUnordered(node.Type) {
+		node.IsUnordered = true
 	}
 
 	for i := 0; i < n.ChildCount(); i++ {
@@ -198,7 +193,7 @@ func buildASTWithRules(n *gotreesitter.Node, src []byte, lang *gotreesitter.Lang
 		}
 	}
 
-	if rules != nil && slices.Contains(rules.Flattened, nodeType) {
+	if rules.IsFlattened(nodeType) {
 		var flattenedChildren []*ASTNode
 		for _, child := range node.Children {
 			flattenedChildren = append(flattenedChildren, child.Children...)
@@ -220,10 +215,6 @@ func unwrapErrorNode(n *gotreesitter.Node, src []byte, lang *gotreesitter.Langua
 	}
 }
 
-func isIgnored(nodeType, label string, ignored []string) bool {
-	return slices.Contains(ignored, nodeType) || slices.Contains(ignored, label)
-}
-
 // Size returns the total number of nodes in the subtree rooted at n.
 func (n *ASTNode) Size() int {
 	if n == nil {

internal/treesitter/rules.go

@@ -4,10 +4,12 @@ import (
 	"embed"
 	"io/fs"
 	"path"
+	"slices"
 
 	"gopkg.in/yaml.v3"
 )
 
+// Rules configures language-specific AST transformations and node matching.
 type Rules struct {
 	Flattened       []string          `yaml:"flattened"`
 	Ignored         []string          `yaml:"ignored"`
@@ -21,16 +23,86 @@ type Rules struct {
 	Pairs           []string          `yaml:"pairs"`
 	Unordered       []string          `yaml:"unordered"`
 	EquivalentTypes [][]string        `yaml:"equivalent_types"`
+
+	flattenedSet    map[string]struct{}
+	ignoredSet      map[string]struct{}
+	labelIgnoredSet map[string]struct{}
+	keywordsSet     map[string]struct{}
+	blocksSet       map[string]struct{}
+	unorderedSet    map[string]struct{}
+	equivGroups     map[string][]int
+}
+
+func (r *Rules) compileSets() {
+	if len(r.Flattened) > 0 {
+		r.flattenedSet = make(map[string]struct{}, len(r.Flattened))
+		for _, s := range r.Flattened {
+			r.flattenedSet[s] = struct{}{}
+		}
+	}
+	if len(r.Ignored) > 0 {
+		r.ignoredSet = make(map[string]struct{}, len(r.Ignored))
+		for _, s := range r.Ignored {
+			r.ignoredSet[s] = struct{}{}
+		}
+	}
+	if len(r.LabelIgnored) > 0 {
+		r.labelIgnoredSet = make(map[string]struct{}, len(r.LabelIgnored))
+		for _, s := range r.LabelIgnored {
+			r.labelIgnoredSet[s] = struct{}{}
+		}
+	}
+	if len(r.Keywords) > 0 {
+		r.keywordsSet = make(map[string]struct{}, len(r.Keywords))
+		for _, s := range r.Keywords {
+			r.keywordsSet[s] = struct{}{}
+		}
+	}
+	if len(r.Blocks) > 0 {
+		r.blocksSet = make(map[string]struct{}, len(r.Blocks))
+		for _, s := range r.Blocks {
+			r.blocksSet[s] = struct{}{}
+		}
+	}
+	if len(r.Unordered) > 0 {
+		r.unorderedSet = make(map[string]struct{}, len(r.Unordered))
+		for _, s := range r.Unordered {
+			r.unorderedSet[s] = struct{}{}
+		}
+	}
+	if len(r.EquivalentTypes) > 0 {
+		r.equivGroups = make(map[string][]int)
+		for idx, group := range r.EquivalentTypes {
+			for _, typ := range group {
+				r.equivGroups[typ] = append(r.equivGroups[typ], idx)
+			}
+		}
+	}
 }
 
-// AreTypesEquivalent checks if t1 and t2 belong to the same equivalent_types group.
+// AreTypesEquivalent checks if t1 and t2 belong to the same equivalence group.
 func (r *Rules) AreTypesEquivalent(t1, t2 string) bool {
 	if r == nil || t1 == "" || t2 == "" {
 		return t1 == t2
 	}
 	if t1 == t2 {
 		return true
 	}
+	if len(r.equivGroups) > 0 {
+		g1, ok1 := r.equivGroups[t1]
+		g2, ok2 := r.equivGroups[t2]
+		if !ok1 || !ok2 {
+			return false
+		}
+		for _, id1 := range g1 {
+			for _, id2 := range g2 {
+				if id1 == id2 {
+					return true
+				}
+			}
+		}
+		return false
+	}
 	for _, group := range r.EquivalentTypes {
 		has1, has2 := false, false
 		for _, typ := range group {
@@ -48,24 +120,126 @@ func (r *Rules) AreTypesEquivalent(t1, t2 string) bool {
 	return false
 }
 
+// IsIgnored checks if a node type or label is filtered out when building the AST.
+func (r *Rules) IsIgnored(nodeType, label string) bool {
+	if r == nil {
+		return false
+	}
+	if len(r.ignoredSet) > 0 {
+		if _, ok := r.ignoredSet[nodeType]; ok {
+			return true
+		}
+		if label != "" {
+			if _, ok := r.ignoredSet[label]; ok {
+				return true
+			}
+		}
+		return false
+	}
+	return slices.Contains(r.Ignored, nodeType) || (label != "" && slices.Contains(r.Ignored, label))
+}
+
+// IsKeyword checks if a node type or label is a language keyword.
+func (r *Rules) IsKeyword(nodeType, label string) bool {
+	if r == nil {
+		return false
+	}
+	if len(r.keywordsSet) > 0 {
+		if _, ok := r.keywordsSet[nodeType]; ok {
+			return true
+		}
+		if label != "" {
+			if _, ok := r.keywordsSet[label]; ok {
+				return true
+			}
+		}
+		return false
+	}
+	return slices.Contains(r.Keywords, nodeType) || (label != "" && slices.Contains(r.Keywords, label))
+}
+
+// IsLabelIgnored checks if node labels should be dropped for this type.
+func (r *Rules) IsLabelIgnored(nodeType string) bool {
+	if r == nil {
+		return false
+	}
+	if len(r.labelIgnoredSet) > 0 {
+		_, ok := r.labelIgnoredSet[nodeType]
+		return ok
+	}
+	return slices.Contains(r.LabelIgnored, nodeType)
+}
+
+// IsUnordered checks if child order doesn't matter for this container.
+func (r *Rules) IsUnordered(nodeType string) bool {
+	if r == nil {
+		return false
+	}
+	if len(r.unorderedSet) > 0 {
+		_, ok := r.unorderedSet[nodeType]
+		return ok
+	}
+	return slices.Contains(r.Unordered, nodeType)
+}
+
+// IsFlattened checks if intermediate nodes of this type should merge into their parent.
+func (r *Rules) IsFlattened(nodeType string) bool {
+	if r == nil {
+		return false
+	}
+	if len(r.flattenedSet) > 0 {
+		_, ok := r.flattenedSet[nodeType]
+		return ok
+	}
+	return slices.Contains(r.Flattened, nodeType)
+}
+
+// IsBlock checks if this node type is a code block.
+func (r *Rules) IsBlock(nodeType string) bool {
+	if r == nil {
+		return false
+	}
+	if len(r.blocksSet) > 0 {
+		_, ok := r.blocksSet[nodeType]
+		return ok
+	}
+	if len(r.Blocks) > 0 {
+		return slices.Contains(r.Blocks, nodeType)
+	}
+	return false
+}
+
+// Alias returns the replacement node type if one exists for the label or node type.
+func (r *Rules) Alias(nodeType, label string) (string, bool) {
+	if r == nil || len(r.Aliased) == 0 {
+		return "", false
+	}
+	if label != "" {
+		if a, ok := r.Aliased[label]; ok {
+			return a, true
+		}
+	}
+	if a, ok := r.Aliased[nodeType]; ok {
+		return a, true
+	}
+	return "", false
+}
+
 //go:embed */rules.yml
 var rulesFS embed.FS
 
-var rulesCache map[string]Rules
+var rulesCache map[string]*Rules
 
+// GetRules returns the compiled AST rules for a language, or nil if none exist.
 func GetRules(lang string) *Rules {
 	if rulesCache == nil {
 		return nil
 	}
-	r, ok := rulesCache[lang]
-	if !ok {
-		return nil
-	}
-	return &r
+	return rulesCache[lang]
 }
 
 func init() {
-	rulesCache = make(map[string]Rules)
+	rulesCache = make(map[string]*Rules)
 	entries, err := fs.ReadDir(rulesFS, ".")
 	if err != nil {
 		panic("failed to read embedded rules directory: " + err.Error())
@@ -82,7 +256,8 @@ func init() {
 			if err := yaml.Unmarshal(data, &r); err != nil {
 				panic("failed to load " + rulePath + ": " + err.Error())
 			}
-			rulesCache[lang] = r
+			r.compileSets()
+			rulesCache[lang] = &r
 		}
 	}
 }