Fix code_lowered_single_expression bug

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

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