From 0704e1c8f3e2b2749df301d95ad51efcf04633bd Mon Sep 17 00:00:00 2001
From: Xianda Sun <sunxdt@gmail.com>
Date: Tue, 9 Apr 2024 06:12:46 +0100
Subject: [PATCH] add recursive strip

---
 src/varname.jl  | 31 +++++++++++++++----------------
 test/varname.jl |  7 +++++++
 2 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/src/varname.jl b/src/varname.jl
index eb58a01..50b2855 100644
--- a/src/varname.jl
+++ b/src/varname.jl
@@ -148,25 +148,24 @@ function Accessors.set(obj, vn::VarName{sym}, value) where {sym}
     return Accessors.set(obj, PropertyLens{sym}() ⨟ getoptic(vn), value)
 end
 
+# recursively stripe off identity
+function stripe_off_identity(optic::ComposedOptic)
+    if optic.inner == identity
+        return stripe_off_identity(optic.outer)
+    elseif optic.outer == identity
+        return stripe_off_identity(optic.inner)
+    else
+        return stripe_off_identity(optic.outer) ∘ stripe_off_identity(optic.inner)
+    end
+end
+stripe_off_identity(optic) = optic
 
 Base.hash(vn::VarName, h::UInt) = hash((getsym(vn), getoptic(vn)), h)
 function Base.:(==)(x::VarName, y::VarName)
     if getsym(x) !== getsym(y)
         return false
     end
-    o1, o2 = map((getoptic(x), getoptic(y))) do o
-        if o isa ComposedOptic
-            if o.inner == identity
-                return o.outer
-            elseif o.outer == identity
-                return o.inner
-            else
-                return o
-            end
-        else
-            return o
-        end
-    end
+    o1, o2 = map(stripe_off_identity, (getoptic(x), getoptic(y)))
     return o1 == o2
 end
 
@@ -647,8 +646,8 @@ function varname(expr::Expr, concretize=Accessors.need_dynamic_optic(expr))
         if concretize
             return :(
                 $(AbstractPPL.VarName){$sym}(
-                    $(AbstractPPL.concretize)($optics, $sym_escaped)
-                )
+                $(AbstractPPL.concretize)($optics, $sym_escaped)
+            )
             )
         elseif Accessors.need_dynamic_optic(expr)
             error("Variable name `$(expr)` is dynamic and requires concretization!")
@@ -703,7 +702,7 @@ function _parse_obj_optics(ex)
         else
             throw(ArgumentError(
                 string("Error while parsing :($ex). Second argument to `getproperty` can only be",
-                       "a `Symbol` or `String` literal, received `$property` instead.")
+                    "a `Symbol` or `String` literal, received `$property` instead.")
             ))
         end
     else
diff --git a/test/varname.jl b/test/varname.jl
index c14c282..8c5cf29 100644
--- a/test/varname.jl
+++ b/test/varname.jl
@@ -56,6 +56,13 @@ end
         vn3 = VarName{:a}(identity ∘ IndexLens(1))
         @test vn1 == vn2
         @test vn1 == vn3
+
+        vn4 = VarName{:a}(identity ∘ PropertyLens{:b}() ∘ (IndexLens(1) ∘ identity))
+        vn5 = VarName{:a}(PropertyLens{:b}() ∘ IndexLens(1))
+        @test vn4 == vn5
+
+        d = Dict(vn4 => 1)
+        @test d[vn5] == 1
     end
 
     @testset "get & set" begin