From 46e843dcda5e6765307af1df1f02ccebde3408e5 Mon Sep 17 00:00:00 2001
From: Charles Kawczynski <kawczynski.charles@gmail.com>
Date: Sat, 9 Mar 2024 16:23:00 -0500
Subject: [PATCH] Fix code_lowered_single_expression bug

---
 src/macro_utils.jl                          | 36 +++++++--------------
 test/expr_code_lowered_single_expression.jl |  2 +-
 2 files changed, 12 insertions(+), 26 deletions(-)

diff --git a/src/macro_utils.jl b/src/macro_utils.jl
index 356194c..44dc752 100644
--- a/src/macro_utils.jl
+++ b/src/macro_utils.jl
@@ -46,32 +46,18 @@ end
 
 fused_pairs(expr::Expr) = Meta.parse(_fused_pairs(expr))
 
-function build_expr(s::String, code_remain)
-    n_subs = count("%", s)
-    if n_subs > 0
-        while n_subs > 0
-            regex = r"%[0-9]"
-            m = match(regex, s)
-            smatch = m.match
-            j = Meta.parse(smatch[2:end])
-            s = replace(s, smatch => string(code_remain[j]))
-            n_subs = count("%", s)
-        end
-    else
-        return s
-    end
-    return s
-end
-
-build_expr(code::Vector) = build_expr(string(code[end]), code)
+# General case: do nothing (identity)
+substitute(x, code) = x
+substitute(x::Core.SSAValue, code) = substitute(code[x.id], code)
+substitute(x::Core.ReturnNode, code) = substitute(code[x.val.id], code)
+substitute(s::Symbol, code) = s
+# Expression: recursively substitute for Expr
+substitute(e::Expr, code) =
+    Expr(substitute(e.head, code), substitute.(e.args, Ref(code))...)
 
+code_info(expr) = Base.Meta.lower(Main, expr).args[1]
 function code_lowered_single_expression(expr)
-    code_lowered = Base.Meta.lower(Main, expr)
-    code_info = code_lowered.args[1]
-    code = code_info.code # vector
-    s = build_expr(code)
-    if startswith(s, "return ")
-        s = replace(s, "return " => "")
-    end
+    code = code_info(expr).code # vector
+    s = string(substitute(code[end], code))
     return Base.Meta.parse(s)
 end
diff --git a/test/expr_code_lowered_single_expression.jl b/test/expr_code_lowered_single_expression.jl
index 15fbfd2..e94cf8e 100644
--- a/test/expr_code_lowered_single_expression.jl
+++ b/test/expr_code_lowered_single_expression.jl
@@ -47,5 +47,5 @@ end
             thermo_params,
         ),
     ))
-    @test_broken MBF.code_lowered_single_expression(expr_in) == expr_out
+    @test MBF.code_lowered_single_expression(expr_in) == expr_out
 end