diff --git a/internal/engine/utils.go b/internal/engine/utils.go index 5379e3f..217764f 100644 --- a/internal/engine/utils.go +++ b/internal/engine/utils.go @@ -37,7 +37,14 @@ func IsTrivialLeaf(n *treesitter.ASTNode) bool { // Returns all descendants in the subtree, excluding n itself. func Descendants(n *treesitter.ASTNode) []*treesitter.ASTNode { - var out []*treesitter.ASTNode + if n == nil { + return nil + } + size := n.Size() + if size <= 1 { + return nil + } + out := make([]*treesitter.ASTNode, 0, size-1) collectDescendants(n, &out) return out } @@ -137,12 +144,14 @@ func PostOrder(n *treesitter.ASTNode) []*treesitter.ASTNode { if n == nil { return nil } - var out []*treesitter.ASTNode + return appendPostOrder(make([]*treesitter.ASTNode, 0, n.Size()), n) +} + +func appendPostOrder(out []*treesitter.ASTNode, n *treesitter.ASTNode) []*treesitter.ASTNode { for _, c := range n.Children { - out = append(out, PostOrder(c)...) + out = appendPostOrder(out, c) } - out = append(out, n) - return out + return append(out, n) } // PreOrder returns all nodes in the subtree rooted at n @@ -152,9 +161,13 @@ func PreOrder(n *treesitter.ASTNode) []*treesitter.ASTNode { if n == nil { return nil } - out := []*treesitter.ASTNode{n} + return appendPreOrder(make([]*treesitter.ASTNode, 0, n.Size()), n) +} + +func appendPreOrder(out []*treesitter.ASTNode, n *treesitter.ASTNode) []*treesitter.ASTNode { + out = append(out, n) for _, c := range n.Children { - out = append(out, PreOrder(c)...) + out = appendPreOrder(out, c) } return out } diff --git a/internal/engine/utils_test.go b/internal/engine/utils_test.go index f2bb023..4a4eb2b 100644 --- a/internal/engine/utils_test.go +++ b/internal/engine/utils_test.go @@ -173,6 +173,9 @@ func TestStructureIsomorphic(t *testing.T) { } func TestPostOrder(t *testing.T) { + if PostOrder(nil) != nil { + t.Error("PostOrder(nil) should return nil") + } leaf := testutil.Leaf("id", "x") root := testutil.Node("call", "", leaf) order := PostOrder(root) @@ -186,6 +189,9 @@ func TestPostOrder(t *testing.T) { } func TestPreOrder(t *testing.T) { + if PreOrder(nil) != nil { + t.Error("PreOrder(nil) should return nil") + } leaf := testutil.Leaf("id", "x") root := testutil.Node("call", "", leaf) order := PreOrder(root)