Support property and index access in ljl_propfunc

legend-exp · Dec 4, 2024 · 9ad126f · 9ad126f
1 parent 3bf5c54
commit 9ad126f
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 13 deletions.
diff --git a/src/ljl_expressions.jl b/src/ljl_expressions.jl
@@ -30,7 +30,7 @@ end
 export parse_ljlexpr
 
 
-const ljl_expr_allowed_heads = (:call, :macrocall, :||, :&&, :comparison)
+const ljl_expr_allowed_heads = [:., :ref, :call, :macrocall, :||, :&&, :comparison]
 
 const ljl_expr_allowed_funcs = Set([
     :!,
@@ -95,17 +95,21 @@ _process_ljlexpr_impl(sym::Symbol, f_varsubst) = f_varsubst(sym)
 function _process_ljlexpr_impl(@nospecialize(expr::Expr), @nospecialize(f_varsubst))
     _process_inner = Base.Fix2(_process_ljlexpr_impl, f_varsubst)
     if expr.head in ljl_expr_allowed_heads
-        if expr.head == :call
+        if expr.head == :.
+            obj = expr.args[begin]
+            propname = expr.args[begin+1]
+            return Expr(expr.head, _process_ljlexpr_impl(obj, f_varsubst), propname)
+        elseif expr.head == :call
             funcname = expr.args[begin]
-            funcargs = expr.args[2:end]
+            funcargs = expr.args[begin+1:end]
             if funcname in ljl_expr_allowed_funcs
                 return Expr(expr.head, funcname, map(_process_inner, funcargs)...)
             else
                 throw(ArgumentError("Function \"$(funcname)\" not allowed in LEGEND Julia expression."))
             end
         elseif expr.head == :macrocall
             macro_name = expr.args[begin]
-            macro_args = expr.args[2:end]
+            macro_args = expr.args[begin+1:end]
             if macro_name == Symbol("@u_str")
                 return Expr(expr.head, macro_name, macro_args...)
             else

diff --git a/test/test_ljl_expressions.jl b/test/test_ljl_expressions.jl
@@ -6,28 +6,31 @@ using Test
 using PropertyFunctions, StructArrays
 import Measurements
 
+using PropDicts: PropDict
+
 include("testing_utils.jl")
 
 @testset "legend_expressions" begin
     data = StructArray([
-        (E_trap = 21092, offs = 0.01, slope = 0.73, A = 328.2, force_accept = false),
-        (E_trap = 21092, offs = 0.01, slope = 0.73, A = NaN, force_accept = false),
-        (E_trap = 21092, offs = 0.01, slope = 0.73, A = NaN, force_accept = true),
+        (E_trap = 21092, offs = 0.01, slope = 0.73, A = 328.2, force_accept = false, a = (b = [1,41,3], c = 51.1)),
+        (E_trap = 21092, offs = 0.01, slope = 0.73, A = NaN, force_accept = false, a = (b = [1,42,3], c = 32.0)),
+        (E_trap = 21092, offs = 0.01, slope = 0.73, A = NaN, force_accept = true, a = (b = [1,41,3], c = 51.1)),
+        (E_trap = 21092, offs = 0.01, slope = 0.73, A = NaN, force_accept = true, a = (b = [1,42,3], c = 32.03)),
     ])
 
-    bool_expr_string = "E_trap > 0 && !isnan(A) && !isinf(E_trap) || force_accept"
+    bool_expr_string = "E_trap > 0 && !isnan(A) && !isinf(E_trap) || force_accept && a.b[2] > a.c"
     bool_expr = parse_ljlexpr(bool_expr_string)
-    @test bool_expr == :(E_trap > 0 && (!(isnan(A)) && !(isinf(E_trap))) || force_accept)
-    ref_boolfunc(x) = (x.E_trap > 0 && !(isnan(x.A)) && !(isinf(x.E_trap)) || x.force_accept)
+    @test bool_expr == :(E_trap > 0 && (!(isnan(A)) && !(isinf(E_trap))) || force_accept && a.b[2] > a.c)
+    ref_boolfunc(x) = (x.E_trap > 0 && !(isnan(x.A)) && !(isinf(x.E_trap)) || x.force_accept && x.a.b[2] > x.a.c)
     bool_pf = ljl_propfunc(bool_expr)
     @test bool_pf isa PropertyFunctions.PropertyFunction
     @test bool_pf === ljl_propfunc(bool_expr_string)
     @test @inferred(broadcast(bool_pf, data)) == ref_boolfunc.(data)
 
-    num_expr_string = "(offs + slope * abs(E_trap) / 1000) - 5.2"
+    num_expr_string = "(offs + slope * abs(E_trap) / 1000) - 5.2 + a.b[2]/100 * a.c/50"
     num_expr = parse_ljlexpr(num_expr_string)
-    @test num_expr == :((offs + (slope * abs(E_trap)) / 1000) - 5.2)
-    ref_numfunc(x) = (x.offs + (x.slope * abs(x.E_trap)) / 1000) - 5.2
+    @test num_expr == :((offs + (slope * abs(E_trap)) / 1000) - 5.2 + a.b[2]/100 * a.c/50)
+    ref_numfunc(x) = (x.offs + (x.slope * abs(x.E_trap)) / 1000) - 5.2 + x.a.b[2]/100 * x.a.c/50
     num_pf = ljl_propfunc(num_expr)
     @test num_pf isa PropertyFunctions.PropertyFunction
     @test num_pf === ljl_propfunc(num_expr_string)