From f76efd431f3739d8cbb37362bb9ae2b187c9a060 Mon Sep 17 00:00:00 2001
From: Paul Berg <paul@plutojl.org>
Date: Sun, 16 Jul 2023 19:18:37 +0200
Subject: [PATCH] ExEx: Track function calls in assignment lhs. (#2604)

---
 src/analysis/ExpressionExplorer.jl | 13 +++++++------
 test/ExpressionExplorer.jl         | 15 +++++++++++++++
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/src/analysis/ExpressionExplorer.jl b/src/analysis/ExpressionExplorer.jl
index 50b773d969..1fd4b5837c 100644
--- a/src/analysis/ExpressionExplorer.jl
+++ b/src/analysis/ExpressionExplorer.jl
@@ -146,7 +146,7 @@ function get_assignees(ex::Expr)::FunctionName
             # e.g. (x, y) in the ex (x, y) = (1, 23)
             args = ex.args
         end
-        union!(Symbol[], Iterators.map(get_assignees, args)...)
+        mapfoldl(get_assignees, union!, args; init=Symbol[])
         # filter(s->s isa Symbol, ex.args)
     elseif ex.head == :(::)
         # TODO: type is referenced
@@ -156,7 +156,7 @@ function get_assignees(ex::Expr)::FunctionName
     elseif ex.head == :...
         # Handles splat assignments. e.g. _, y... = 1:5
         args = ex.args
-        union!(Symbol[], Iterators.map(get_assignees, args)...)
+        mapfoldl(get_assignees, union!, args; init=Symbol[])
     else
         @warn "unknown use of `=`. Assignee is unrecognised." ex
         Symbol[]
@@ -373,12 +373,13 @@ function explore_assignment!(ex::Expr, scopestate::ScopeState)::SymbolsState
     global_assignees = get_global_assignees(assignees, scopestate)
 
     # If we are _not_ assigning a global variable, then this symbol hides any global definition with that name
-    push!(scopestate.hiddenglobals, setdiff(assignees, global_assignees)...)
+    union!(scopestate.hiddenglobals, setdiff(assignees, global_assignees))
     assigneesymstate = explore!(ex.args[1], scopestate)
 
-    push!(scopestate.hiddenglobals, global_assignees...)
-    push!(symstate.assignments, global_assignees...)
-    push!(symstate.references, setdiff(assigneesymstate.references, global_assignees)...)
+    union!(scopestate.hiddenglobals, global_assignees)
+    union!(symstate.assignments, global_assignees)
+    union!(symstate.references, setdiff(assigneesymstate.references, global_assignees))
+    union!(symstate.funccalls, filter!(call -> length(call) != 1 || only(call) ∉ global_assignees, assigneesymstate.funccalls))
     filter!(!all_underscores, symstate.references)  # Never record _ as a reference
 
     return symstate
diff --git a/test/ExpressionExplorer.jl b/test/ExpressionExplorer.jl
index 505b7470ea..afe010c1ff 100644
--- a/test/ExpressionExplorer.jl
+++ b/test/ExpressionExplorer.jl
@@ -165,6 +165,9 @@ Some of these @test_broken lines are commented out to prevent printing to the te
         @test testee(:(x = let a = 1; a += b end), [:b], [:x], [:+], [])
         @test testee(:(_ = a + 1), [:a], [], [:+], [])
         @test testee(:(a = _ + 1), [], [:a], [:+], [])
+
+        @test testee(:(f()[] = 1), [], [], [:f], [])
+        @test testee(:(x[f()] = 1), [:x], [], [:f], [])
     end
     @testset "Multiple assignments" begin
         # Note that using the shorthand syntax :(a = 1, b = 2) to create an expression
@@ -206,6 +209,18 @@ Some of these @test_broken lines are commented out to prevent printing to the te
         @test testee(quote
             a, b... = 0:5
         end, [],[:a, :b], [[:(:)]], [])
+        @test testee(quote
+            a[x], x = 1, 2
+        end, [:a], [:x], [], [])
+        @test testee(quote
+            x, a[x] = 1, 2
+        end, [:a], [:x], [], [])
+        @test testee(quote
+            f, a[f()] = g
+        end, [:g, :a], [:f], [], [])
+        @test testee(quote
+            a[f()], f = g
+        end, [:g, :a], [:f], [], [])
         @test testee(quote (; a, b) = x end, [:x], [:a, :b], [], [])
         @test testee(quote a = (b, c) end, [:b, :c], [:a], [], [])