compiler: stamp node ids once, at compile time

lib/lua/ast/ids.ex

@@ -8,8 +8,8 @@ defmodule Lua.AST.Ids do
   node's whole subtree, so the deepest nodes (function bodies, blocks) — the
   ones that make the most useful keys — are also the most expensive ones.
 
-  `assign/1` walks a parsed chunk once and writes a distinct integer into
-  each node's `meta.id`, letting those tables key on a single word instead.
+  `assign/1` walks a chunk once and writes a distinct integer into each
+  node's `meta.id`, letting those tables key on a single word instead.
   Ids are unique across the whole chunk, including the bodies of nested
   functions.
   """
@@ -23,8 +23,10 @@ defmodule Lua.AST.Ids do
   @doc """
   Returns `chunk` with every reachable node stamped with a unique `meta.id`.
 
-  Nodes the parser built without metadata gain a `Lua.AST.Meta` carrying only
-  the id; nodes that already have one keep their positions and comments.
+  Nodes built without metadata gain a `Lua.AST.Meta` carrying only the id;
+  nodes that already have one keep their positions and comments. Ids are
+  reassigned from scratch on every call, so a partially stamped chunk comes
+  back fully and consistently numbered.
   """
   @spec assign(Chunk.t()) :: Chunk.t()
   def assign(%Chunk{} = chunk) do

lib/lua/compiler.ex

@@ -32,10 +32,12 @@ defmodule Lua.Compiler do
   def compile(%Chunk{} = chunk, opts \\ []) do
     # Scope resolution and codegen key per-node tables by `meta.id` (see
     # `Lua.AST.Ids`); without ids, structurally identical nodes (e.g. two
-    # empty loop bodies) would share one table entry and miscompile. Stamping
-    # here covers chunks that never went through the parser, such as those
-    # built with `Lua.AST.Builder`. Assignment is deterministic, so a parsed
-    # chunk (already stamped by `Lua.Parser`) re-stamps to the same ids.
+    # empty loop bodies) would share one table entry and miscompile. This is
+    # the single stamping point, so it covers parsed chunks and chunks built
+    # by hand (`Lua.AST.Builder`) alike. It cannot be skipped for an
+    # already-stamped chunk: a chunk carrying hand-built nodes spliced into a
+    # parsed one is partially stamped, and proving otherwise costs the walk
+    # this would save.
     chunk = Ids.assign(chunk)
 
     with :ok <- GotoValidation.validate(chunk),

lib/lua/parser.ex

@@ -8,7 +8,6 @@ defmodule Lua.Parser do
   alias Lua.AST.Block
   alias Lua.AST.Chunk
   alias Lua.AST.Expr
-  alias Lua.AST.Ids
   alias Lua.AST.Meta
   alias Lua.AST.Statement
   alias Lua.Lexer
@@ -107,16 +106,16 @@ defmodule Lua.Parser do
   @doc """
   Parses a chunk (top-level block) from a token list.
 
-  Every node of the returned chunk carries a chunk-unique `meta.id`; see
-  `Lua.AST.Ids`.
+  Node ids are not stamped here; `Lua.Compiler.compile/2` stamps every chunk
+  it compiles, including chunks built without the parser. See `Lua.AST.Ids`.
   """
   @spec parse_chunk([token()]) :: {:ok, Chunk.t()} | {:error, term()}
   def parse_chunk(tokens) do
     case parse_block(tokens) do
       {:ok, block, rest} ->
         case rest do
           [{:eof, _}] ->
-            {:ok, Ids.assign(Chunk.new(block))}
+            {:ok, Chunk.new(block)}
 
           [{type, _, pos} | _] ->
             {:error, {:unexpected_token, type, pos, "Expected end of input"}}

test/lua/ast/ids_test.exs

@@ -54,7 +54,7 @@ defmodule Lua.AST.IdsTest do
 
     test "keeps positions and comments already on a node" do
       {:ok, chunk} = Lua.Parser.parse_raw("-- leading\nlocal x = 1\n")
-      [local_stmt] = chunk.block.stmts
+      [local_stmt] = Ids.assign(chunk).block.stmts
 
       assert %{line: 2} = local_stmt.meta.start
       assert [%{text: " leading"}] = local_stmt.meta.metadata.leading_comments
@@ -63,8 +63,9 @@ defmodule Lua.AST.IdsTest do
 
     test "is idempotent in shape: re-assigning yields the same chunk" do
       {:ok, chunk} = Lua.Parser.parse_raw("local t = {1, 2, x = 3}\nreturn t.x\n")
+      stamped = Ids.assign(chunk)
 
-      assert Ids.assign(chunk) == chunk
+      assert Ids.assign(stamped) == stamped
     end
 
     test "numbers every node of the compilable surface, uniquely" do
@@ -90,6 +91,8 @@ defmodule Lua.AST.IdsTest do
   defp ids_for(source) do
     {:ok, chunk} = Lua.Parser.parse_raw(source)
 
-    Walker.reduce(chunk, [], fn node, acc -> [node.meta.id | acc] end)
+    chunk
+    |> Ids.assign()
+    |> Walker.reduce([], fn node, acc -> [node.meta.id | acc] end)
   end
 end