From ef20234e6792c3a40a304496f1c116541ca47d4f Mon Sep 17 00:00:00 2001
From: arichard <richard4@cs.stanford.edu>
Date: Tue, 10 May 2022 10:25:03 -0700
Subject: [PATCH] Also stops tracking at the *end* of control flow blocks

---
 spec/table_field_limitations_spec.lua      | 27 ++++++++++++++++++++++
 src/luacheck/stages/check_table_fields.lua | 11 ++++++++-
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/spec/table_field_limitations_spec.lua b/spec/table_field_limitations_spec.lua
index 84aa865b..c148e596 100644
--- a/spec/table_field_limitations_spec.lua
+++ b/spec/table_field_limitations_spec.lua
@@ -130,6 +130,33 @@ y[1] = 1
       ]])
    end)
 
+   it("halts checking at the end of control flow blocks with jumps", function()
+      assert_warnings({}, [[
+local x = {1}
+if math.rand(0,1) ~= 1 then
+   x = {}
+end
+
+x[1] = x[1]
+
+local y = {1}
+if math.random(0,1) == 1 then
+   y[1] = 2
+else
+   y = {}
+end
+
+y[1] = y[1]
+
+local a = {1}
+while math.random(0,1) == 1 do
+   a = {}
+end
+
+a[1] = a[1]
+      ]])
+   end)
+
    it("stops checking if a function is called", function()
       assert_warnings({
          {line = 8, column = 3, name = 'y', end_column = 3, field = 'x', code = '315', set_is_nil = '' },
diff --git a/src/luacheck/stages/check_table_fields.lua b/src/luacheck/stages/check_table_fields.lua
index 6db41a24..9e36e6d6 100644
--- a/src/luacheck/stages/check_table_fields.lua
+++ b/src/luacheck/stages/check_table_fields.lua
@@ -25,6 +25,11 @@ function ClosureState:__init(chstate)
    self.external_references_accessed = {}
 
    self.max_item = 0
+
+   -- Luacheck's linearized item format doesn't explicitly have a node for the "end" tag
+   -- This tracks where end tags would be, so that we can stop tracking
+   -- This pattern won't be workable if this eventually gets extended to support control flow
+   self.jump_destinations = {}
 end
 
 -- Start keeping track of a local table
@@ -417,6 +422,7 @@ function ClosureState:handle_eval(item)
 end
 
 function ClosureState:handle_jump(item)
+   self.jump_destinations[item.to] = true
    if item.to > self.max_item then
       -- return; see comment under handle_control_flow_item
       self:on_scope_end()
@@ -428,7 +434,7 @@ end
 local item_callbacks = {
    Noop = ClosureState.handle_control_flow_item,
    Jump = ClosureState.handle_jump,
-   Cjump = ClosureState.stop_tracking_tables,
+   Cjump = ClosureState.handle_jump,
    Eval = ClosureState.handle_eval,
    Local = ClosureState.handle_local_or_set_item,
    Set = ClosureState.handle_local_or_set_item
@@ -453,6 +459,9 @@ local function detect_unused_table_fields(closure, check_state)
    closure_state.max_item = #closure.items
 
    for item_index = 1, #closure.items do
+      if closure_state.jump_destinations[item_index] then
+         closure_state:stop_tracking_tables()
+      end
       -- Function declaration: function could potentially survive this scope
       -- Preserving a reference to its upvalues
       local item = closure.items[item_index]