Skip to content

Commit 525370a

Browse files
committed
refactor: better printing for assignment operator
1 parent 1ff3bba commit 525370a

File tree

3 files changed

+33
-12
lines changed

3 files changed

+33
-12
lines changed

src/SpecialOperators.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ Base.@kwdef struct AssignOperator <: Function
2323
end
2424
@declare_expression_operator((op::AssignOperator), 1)
2525
@inline special_operator(::Type{AssignOperator}) = true
26-
get_op_name(o::AssignOperator) = "[{FEATURE_" * string(o.target_register) * "} =]"
26+
get_op_name(o::AssignOperator) = "ASSIGN_OP:{FEATURE_" * string(o.target_register) * "}"
2727

2828
function deg1_eval_special(tree, cX, operators, op::AssignOperator, eval_options)
2929
result = _eval_tree_array(tree.l, cX, operators, eval_options)

src/Strings.jl

Lines changed: 27 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -116,12 +116,33 @@ function combine_op_with_inputs(op, l, r)::Vector{Char}
116116
end
117117
end
118118
function combine_op_with_inputs(op, l)
119-
# "op(l)"
120-
out = copy(op)
121-
push!(out, '(')
122-
append!(out, strip_brackets(l))
123-
push!(out, ')')
124-
return out
119+
# Check if this is an assignment operator with our special prefix
120+
op_str = String(op)
121+
if startswith(op_str, "ASSIGN_OP:")
122+
# Extract the variable name from the operator name
123+
var_name = op_str[11:end]
124+
# Format: (var ← expr)
125+
out = ['(']
126+
append!(out, collect(var_name))
127+
append!(out, collect(""))
128+
# Ensure the expression is always wrapped in parentheses for clarity
129+
if l[1] == '(' && l[end] == ')'
130+
append!(out, l)
131+
else
132+
push!(out, '(')
133+
append!(out, strip_brackets(l))
134+
push!(out, ')')
135+
end
136+
push!(out, ')')
137+
return out
138+
else
139+
# Regular unary operator: "op(l)"
140+
out = copy(op)
141+
push!(out, '(')
142+
append!(out, strip_brackets(l))
143+
push!(out, ')')
144+
return out
145+
end
125146
end
126147

127148
"""

test/test_special_operators.jl

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ using TestItems: @testitem
2323
x2 = Expression(Node(; feature=2); operators, variable_names)
2424
assign_expr = assign_x2(0.0 * x1 + 3.0) + x2
2525

26-
@test string_tree(assign_expr) == "[x2 =]((0.0 * x1) + 3.0) + x2"
26+
@test string_tree(assign_expr) == "(x2 ((0.0 * x1) + 3.0)) + x2"
2727

2828
# We should see that x2 will become 3.0 _before_ adding
2929
result, completed = eval_tree_array(assign_expr, X)
@@ -35,7 +35,7 @@ using TestItems: @testitem
3535

3636
# But, with the reverse order, we get the x2 _before_ it was reassigned
3737
assign_expr_reverse = x2 + assign_x2(0.0 * x1 + 3.0)
38-
@test string_tree(assign_expr_reverse) == "x2 + [x2 =]((0.0 * x1) + 3.0)"
38+
@test string_tree(assign_expr_reverse) == "x2 + (x2 ((0.0 * x1) + 3.0))"
3939
result, completed = eval_tree_array(assign_expr_reverse, X)
4040
@test completed == true
4141
@test result == [3.5, 4.5, 5.5]
@@ -58,7 +58,7 @@ end
5858
x3 = Expression(Node(; feature=3); operators, variable_names)
5959

6060
expr = assign_x1(assign_x1(x1 * 2) + x1)
61-
@test string_tree(expr) == "[a =]([a =](a * 2.0) + a)"
61+
@test string_tree(expr) == "a ← ((a ← (a * 2.0)) + a)"
6262

6363
result, completed = eval_tree_array(expr, X)
6464
@test completed == true
@@ -122,7 +122,7 @@ end
122122
x2 = Expression(Node(; feature=2); operators, variable_names)
123123
expr = while_op(3.0 - x2, assign_x2(x2 + 1.0))
124124

125-
@test string_tree(expr) == "while(3.0 - x2, [x2 =](x2 + 1.0))"
125+
@test string_tree(expr) == "while(3.0 - x2, x2 (x2 + 1.0))"
126126

127127
result, completed = eval_tree_array(expr, X)
128128
@test completed == true
@@ -163,7 +163,7 @@ end
163163
expr = (while_op(condition, body) * 0.0) + xs[3]
164164

165165
@test string_tree(expr) ==
166-
"(while(x2, (([x5 =](x3) + [x3 =](x4)) + [x4 =](x5 + x4)) + [x2 =](x2 - 1.0)) * 0.0) + x3"
166+
"(while(x2, (((x5 (x3)) + (x3 (x4))) + (x4 (x5 + x4))) + (x2 (x2 - 1.0))) * 0.0) + x3"
167167

168168
result, completed = eval_tree_array(expr, X)
169169
@test completed == true

0 commit comments

Comments
 (0)