Always shadow for-loop control variables in macro bodies

spec/lang/macro/for_control_var_spec.lua

@@ -0,0 +1,66 @@
+local tl = require('tl')
+local lua_generator = require('teal.gen.lua_generator')
+
+describe('macro body for-loop control variables', function()
+   it('allows a numeric for body to write to the control variable', function()
+      local code = [[
+         local macro thrice!(x: Expression)
+            local out = block('statements')
+            for i = 1, 3 do
+               i = i + 10
+               table.insert(out, `$x`)
+            end
+            return out
+         end
+
+         thrice!(print('hi'))
+      ]]
+      local ast, errs = tl.parse(code)
+      assert.same({}, errs)
+      local out, err = lua_generator.generate(ast, '5.4')
+      assert.is_nil(err)
+      out = out:gsub("^%s+", ""):gsub("%s+$", "")
+      assert.same("print('hi'); print('hi'); print('hi')", out)
+   end)
+
+   it('allows a generic for body to write to the control variable', function()
+      local code = [[
+         local macro each!(x: Expression)
+            local out = block('statements')
+            for word in ('a b'):gmatch('%a+') do
+               word = word .. '!'
+               table.insert(out, `$x`)
+            end
+            return out
+         end
+
+         each!(print('hi'))
+      ]]
+      local ast, errs = tl.parse(code)
+      assert.same({}, errs)
+      local out, err = lua_generator.generate(ast, '5.4')
+      assert.is_nil(err)
+      out = out:gsub("^%s+", ""):gsub("%s+$", "")
+      assert.same("print('hi'); print('hi')", out)
+   end)
+
+   it('does not disturb loops that never write the control variable', function()
+      local code = [[
+         local macro twice!(x: Expression)
+            local out = block('statements')
+            for _ = 1, 2 do
+               table.insert(out, `$x`)
+            end
+            return out
+         end
+
+         twice!(print('hi'))
+      ]]
+      local ast, errs = tl.parse(code)
+      assert.same({}, errs)
+      local out, err = lua_generator.generate(ast, '5.4')
+      assert.is_nil(err)
+      out = out:gsub("^%s+", ""):gsub("%s+$", "")
+      assert.same("print('hi'); print('hi')", out)
+   end)
+end)

teal.lua

@@ -13760,6 +13760,40 @@ local function is_statement_kind(k)
    k == "record_function" or k == "newtype" or k == "pragma"
 end
 
+
+local function add_for_shadows(b)
+   if b.kind == "forin" or b.kind == "fornum" then
+      local map = BLOCK_INDEXES[b.kind:upper()]
+      local body = b[map.BODY]
+
+      local var = b[1]
+      if b.kind == "forin" then
+         var = var and var[1]
+      end
+
+      if body and var and var.kind == "identifier" then
+         local name = var.tk
+
+         local function at(kind)
+            return { kind = kind, tk = name, f = b.f, y = b.y, x = b.x, yend = b.y, xend = b.x }
+         end
+
+         local LD = BLOCK_INDEXES.LOCAL_DECLARATION
+         local stmt = at("local_declaration")
+         stmt[LD.VARS] = at("variable_list")
+         stmt[LD.VARS][1] = at("identifier")
+         stmt[LD.EXPS] = at("expression_list")
+         stmt[LD.EXPS][1] = at("identifier")
+
+         table.insert(body, 1, stmt)
+      end
+   end
+
+   for _, child in children(b) do
+      add_for_shadows(child)
+   end
+end
+
 local function compile_local_macro(mb, filename, read_lang, env, errs)
    local name_block = mb[BLOCK_INDEXES.LOCAL_MACRO.NAME]
    if not name_block or name_block.kind ~= "identifier" then
@@ -13802,6 +13836,9 @@ local function compile_local_macro(mb, filename, read_lang, env, errs)
    local lua_generator = require("teal.gen.lua_generator")
    local single = { kind = "statements", y = mb.y, x = mb.x, tk = mb.tk, yend = mb.yend, xend = mb.xend }
    single[1] = mb
+
+   add_for_shadows(single)
+
    local mast, perrs = ast.parse_blocks(single, filename, read_lang)
    if #perrs > 0 then
       for _, e in ipairs(perrs) do table.insert(errs, e) end

teal/macro_eval.lua

@@ -178,6 +178,40 @@ local function is_statement_kind(k)
    k == "record_function" or k == "newtype" or k == "pragma"
 end
 
+
+local function add_for_shadows(b)
+   if b.kind == "forin" or b.kind == "fornum" then
+      local map = BLOCK_INDEXES[b.kind:upper()]
+      local body = b[map.BODY]
+
+      local var = b[1]
+      if b.kind == "forin" then
+         var = var and var[1]
+      end
+
+      if body and var and var.kind == "identifier" then
+         local name = var.tk
+
+         local function at(kind)
+            return { kind = kind, tk = name, f = b.f, y = b.y, x = b.x, yend = b.y, xend = b.x }
+         end
+
+         local LD = BLOCK_INDEXES.LOCAL_DECLARATION
+         local stmt = at("local_declaration")
+         stmt[LD.VARS] = at("variable_list")
+         stmt[LD.VARS][1] = at("identifier")
+         stmt[LD.EXPS] = at("expression_list")
+         stmt[LD.EXPS][1] = at("identifier")
+
+         table.insert(body, 1, stmt)
+      end
+   end
+
+   for _, child in children(b) do
+      add_for_shadows(child)
+   end
+end
+
 local function compile_local_macro(mb, filename, read_lang, env, errs)
    local name_block = mb[BLOCK_INDEXES.LOCAL_MACRO.NAME]
    if not name_block or name_block.kind ~= "identifier" then
@@ -220,6 +254,9 @@ local function compile_local_macro(mb, filename, read_lang, env, errs)
    local lua_generator = require("teal.gen.lua_generator")
    local single = { kind = "statements", y = mb.y, x = mb.x, tk = mb.tk, yend = mb.yend, xend = mb.xend }
    single[1] = mb
+
+   add_for_shadows(single)
+
    local mast, perrs = ast.parse_blocks(single, filename, read_lang)
    if #perrs > 0 then
       for _, e in ipairs(perrs) do table.insert(errs, e) end

teal/macro_eval.tl

@@ -178,6 +178,40 @@ local function is_statement_kind(k: block.BlockKind): boolean
           k == "record_function" or k == "newtype" or k == "pragma"
 end
 
+-- 5.5 makes `for` control variables const; shadowing is a no-op in older targets.
+local function add_for_shadows(b: Block)
+   if b.kind == "forin" or b.kind == "fornum" then
+      local map = BLOCK_INDEXES[b.kind:upper()]
+      local body = b[map.BODY]
+
+      local var = b[1]
+      if b.kind == "forin" then
+         var = var and var[1]
+      end
+
+      if body and var and var.kind == "identifier" then
+         local name = var.tk
+
+         local function at(kind : block.BlockKind) : Block
+            return { kind = kind, tk = name, f = b.f, y = b.y, x = b.x, yend = b.y, xend = b.x }
+         end
+
+         local LD = BLOCK_INDEXES.LOCAL_DECLARATION
+         local stmt = at("local_declaration")
+         stmt[LD.VARS] = at("variable_list")
+         stmt[LD.VARS][1] = at("identifier")
+         stmt[LD.EXPS] = at("expression_list") 
+         stmt[LD.EXPS][1] = at("identifier")
+
+         table.insert(body, 1, stmt)
+      end
+   end
+
+   for _, child in children(b) do
+      add_for_shadows(child)
+   end
+end
+
 local function compile_local_macro(mb: Block, filename: string, read_lang: BlockLang, env: MacroEnv, errs: {Error})
    local name_block = mb[BLOCK_INDEXES.LOCAL_MACRO.NAME]
    if not name_block or name_block.kind ~= "identifier" then
@@ -220,6 +254,9 @@ local function compile_local_macro(mb: Block, filename: string, read_lang: Block
    local lua_generator = require("teal.gen.lua_generator")
    local single: Block = { kind = "statements", y = mb.y, x = mb.x, tk = mb.tk, yend = mb.yend, xend = mb.xend }
    single[1] = mb
+
+   add_for_shadows(single)
+
    local mast, perrs = ast.parse_blocks(single, filename, read_lang)
    if #perrs > 0 then
       for _, e in ipairs(perrs) do table.insert(errs, e) end

tl.lua

@@ -13604,6 +13604,40 @@ local function is_statement_kind(k)
    k == "record_function" or k == "newtype" or k == "pragma"
 end
 
+
+local function add_for_shadows(b)
+   if b.kind == "forin" or b.kind == "fornum" then
+      local map = BLOCK_INDEXES[b.kind:upper()]
+      local body = b[map.BODY]
+
+      local var = b[1]
+      if b.kind == "forin" then
+         var = var and var[1]
+      end
+
+      if body and var and var.kind == "identifier" then
+         local name = var.tk
+
+         local function at(kind)
+            return { kind = kind, tk = name, f = b.f, y = b.y, x = b.x, yend = b.y, xend = b.x }
+         end
+
+         local LD = BLOCK_INDEXES.LOCAL_DECLARATION
+         local stmt = at("local_declaration")
+         stmt[LD.VARS] = at("variable_list")
+         stmt[LD.VARS][1] = at("identifier")
+         stmt[LD.EXPS] = at("expression_list")
+         stmt[LD.EXPS][1] = at("identifier")
+
+         table.insert(body, 1, stmt)
+      end
+   end
+
+   for _, child in children(b) do
+      add_for_shadows(child)
+   end
+end
+
 local function compile_local_macro(mb, filename, read_lang, env, errs)
    local name_block = mb[BLOCK_INDEXES.LOCAL_MACRO.NAME]
    if not name_block or name_block.kind ~= "identifier" then
@@ -13646,6 +13680,9 @@ local function compile_local_macro(mb, filename, read_lang, env, errs)
    local lua_generator = require("teal.gen.lua_generator")
    local single = { kind = "statements", y = mb.y, x = mb.x, tk = mb.tk, yend = mb.yend, xend = mb.xend }
    single[1] = mb
+
+   add_for_shadows(single)
+
    local mast, perrs = ast.parse_blocks(single, filename, read_lang)
    if #perrs > 0 then
       for _, e in ipairs(perrs) do table.insert(errs, e) end