Same wildcard should match the same subtree

fixes #15

Author: cshuaimin

lua/ssr/search.lua

@@ -41,6 +41,33 @@ function ExtmarkRange:get()
   return extmark[1], extmark[2], extmark[3].end_row, extmark[3].end_col
 end
 
+ts.query.add_predicate("is-same-tree?", function(match, _pattern, buf, pred)
+  ---@param node1 TSNode
+  ---@param node2 TSNode
+  ---@return boolean
+  local function is_same_tree(node1, node2)
+    if node1:type() ~= node2:type() then
+      return false
+    end
+    if node1:child_count() ~= node2:child_count() then
+      return false
+    end
+    if node1:child_count() == 0 then
+      if ts.get_node_text(node1, buf) ~= ts.get_node_text(node2, buf) then
+        return false
+      end
+    else
+      for i = 0, node1:child_count() - 1 do
+        if not is_same_tree(node1:child(i), node2:child(i)) then
+          return false
+        end
+      end
+    end
+    return true
+  end
+  return is_same_tree(match[pred[2]], match[pred[3]])
+end, true)
+
 -- Build a TS sexpr represting the node.
 ---@param node TSNode
 ---@param source string
@@ -56,9 +83,16 @@ local function build_sexpr(node, source)
     -- Special identifier __ssr_var_name is a named wildcard.
     local var = text:match("^" .. wildcard_prefix .. "([_%a%d]+)$")
     if var then
-      wildcards[var] = next_idx
-      next_idx = next_idx + 1
-      return "(_) @" .. var
+      if not wildcards[var] then
+        wildcards[var] = next_idx
+        next_idx = next_idx + 1
+        return "(_) @" .. var
+      else
+        -- Same wildcard should match the same subtree.
+        local sexpr = string.format("(_) @_%d (#is-same-tree? @_%d @%s)", next_idx, next_idx, var)
+        next_idx = next_idx + 1
+        return sexpr
+      end
     end
 
     -- Leaf nodes (keyword, identifier, literal and symbol) should match text.

tests/ssr_spec.lua

@@ -141,7 +141,6 @@ String::from((y + 5).foo(z))
 t [[ go parsed correctly
 func main() {
   <commit, _ := os.LookupEnv("GITHUB_SHA")>
-  print(commit)
 }
 ====
 $a, _ := os.LookupEnv($b)
@@ -150,10 +149,43 @@ $a := os.Getenv($b)
 ====
 func main() {
   commit := os.Getenv("GITHUB_SHA")
-  print(commit)
 }
 ]]
 
+t [[ rust match same tree
+<idx = idx + 1>;
+bar = foo + idx;
+*foo.bar() = * foo . bar () + 1;
+(foo + bar) = (foo + bar) + 1;
+(foo + bar) = (foo - bar) + 1;
+====
+$a = $a + $b ==>> $a += $b
+====
+idx += 1;
+bar = foo + idx;
+*foo.bar() += 1;
+(foo + bar) += 1;
+(foo + bar) = (foo - bar) + 1;
+]]
+
+t [[ python correct indent for same tree
+def f():
+    <if await foo.bar(baz):
+        if await foo.bar(baz):
+            pass>
+====
+if $foo:
+    if $foo:
+        $body
+==>>
+if $foo:
+    $body
+====
+def f():
+    if await foo.bar(baz):
+        pass
+]]
+
 describe("", function()
   for _, s in ipairs(tests) do
     local ft, desc, content, pattern, template, expected =