working but slow

src/optics.jl

@@ -42,8 +42,7 @@ julia> obj = (a=1, b=2); lens=@optic _.a; val = 100;
 
 julia> set(obj, lens, val)
 (a = 100, b = 2)
-```
-See also [`modify`](@ref).
+``` See also [`modify`](@ref).
 """
 function set end
 
@@ -346,15 +345,7 @@ Here `f` has signature `f(::Value, ::State) -> Tuple{NewValue, NewState}`.
 """
 function modify_stateful end
 
-@inline function modify_stateful(f, (obj, state), optic::Properties)
- let f=f, obj=obj, state=state
- modify_stateful_context((obj, state), optic) do _, fn, pr, st
- f(getfield(pr, known(fn)), st)
- end
- end
-end
-
-@generated function modify_stateful_context(f, (obj, state1)::T, optic::Properties) where T
+@generated function modify_stateful(f::F, (obj, state)::T, optic::Properties) where {T,F}
  _modify_stateful_inner(T)
 end
 
@@ -363,29 +354,29 @@ function _modify_stateful_inner(::Type{<:Tuple{O,S}}) where {O,S}
  modifications = []
  vals = Expr(:tuple)
  fns = fieldnames(O)
- local st1 = :state0
- local st2 = :state1
  for (i, fn) in enumerate(fns)
  v = Symbol("val$i")
- st1 = Symbol("state$i")
- st2 = Symbol("state$(i+1)")
- ms = if O <: Tuple 
- :(($v, $st2) = f(obj, StaticInt{$(QuoteNode(fn))}(), props, $st1))
+ st = if S <: ContextState 
+ if O <: Tuple 
+ :(ContextState(state.vals, obj, StaticInt{$(QuoteNode(fn))}()))
+ else
+ :(ContextState(state.vals, obj, StaticSymbol{$(QuoteNode(fn))}()))
+ end
  else
- :(($v, $st2) = f(obj, StaticSymbol{$(QuoteNode(fn))}(), props, $st1))
+ :state
  end
+ ms = :(($v, state) = f(getfield(props, $(QuoteNode(fn))), $st))
  push!(modifications, ms)
  push!(vals.args, v)
  end
  patch = O <: Tuple ? vals : :(NamedTuple{$fns}($vals))
- Expr(:block,
- :(props = getproperties(obj)),
- modifications...,
- :(patch = $patch),
- :(new_obj = maybesetproperties($st2, obj, patch)),
- :(new_state = maybesetstate($st2, obj, patch)),
- :(return (setproperties(obj, patch), $st2)),
- )
+ start = :(props = getproperties(obj))
+ rest = MacroTools.@q begin
+ patch = $patch
+ new_obj = maybesetproperties(state, obj, patch)
+ return (new_obj, state)
+ end
+ Expr(:block, start, modifications..., rest) 
 end
 
 maybesetproperties(state, obj, patch) = setproperties(obj, patch)
@@ -426,15 +417,10 @@ Query(; select=Any, descend=x -> true, optic=Properties()) = Query(select, desce
 
 OpticStyle(::Type{<:AbstractQuery}) = SetBased()
 
-struct Context{Select,Descend,Optic<:Union{ComposedOptic,Properties}} <: AbstractQuery
- select_condition::Select
- descent_condition::Descend
- optic::Optic
-end
-
-
-struct ContextState{V}
+struct ContextState{V,O,FN}
  vals::V
+ obj::O
+ fn::FN
 end
 struct GetAllState{V}
  vals::V
@@ -445,57 +431,69 @@ struct SetAllState{C,V,I}
  itr::I
 end
 
-pop(x) = first(x), Base.tail(x)
-push(x, val) = (x..., val)
-push(x::GetAllState, val) = GetAllState(push(x.vals, val))
+const GetStates = Union{GetAllState,ContextState}
+
+@inline pop(x) = first(x), Base.tail(x)
+@inline push(x, val) = (x..., val)
+@inline push(x::GetAllState, val) = GetAllState(push(x.vals, val))
+@inline push(x::ContextState, val) = ContextState(push(x.vals, val), nothing, nothing)
 
 (q::Query)(obj) = getall(obj, q)
 
-function getall(obj, q)
+getall(obj, q) = _getall(obj, q).vals
+function _getall(obj, q::Q) where Q<:Query
  initial_state = GetAllState(())
- _, final_state = modify_stateful((obj, initial_state), q) do o, s
- new_state = push(s, outer(q.optic, o, s))
- o, new_state
+ _, final_state = let q=q
+ modify_stateful((obj, initial_state), q) do o, s
+ new_state = push(s, outer(q.optic, o, s))
+ o, new_state
+ end
  end
- return final_state.vals
+ final_state
 end
 
-function setall(obj, q, vals)
+function setall(obj, q::Q, vals) where Q<:Query
  initial_state = SetAllState(Unchanged(), vals, 1)
- final_obj, _ = modify_stateful((obj, initial_state), q) do o, s
- new_output = outer(q.optic, o, s)
- new_state = SetAllState(Changed(), s.vals, s.itr + 1)
- new_output, new_state
+ final_obj, _ = let obj=obj, q=q, initial_state=initial_state
+ modify_stateful((obj, initial_state), q) do o, s
+ new_output = outer(q.optic, o, s)
+ new_state = SetAllState(Changed(), s.vals, s.itr + 1)
+ new_output, new_state
+ end
  end
  return final_obj
 end
 
-function context(f, obj, q)
- initial_state = GetAllState(())
- _, final_state = modify_stateful_context((obj, initial_state), Properties()) do o, fn, pr, s
- new_state = push(s, f(o, known(fn)))
- o, new_state
+function context(f::F, obj, q::Q) where {F,Q<:Query}
+ initial_state = ContextState((), nothing, nothing)
+ _, final_state = let f=f
+ modify_stateful((obj, initial_state), q) do o, s
+ new_state = push(s, f(s.obj, known(s.fn)))
+ o, new_state
+ end
  end
  return final_state.vals
 end
 
 modify(f, obj, q::Query) = setall(obj, q, map(f, getall(obj, q)))
 
-@inline function modify_stateful(f::F, (obj, state), q::Query) where F
- modify_stateful((obj, state), inner(q.optic)) do o, s
- if q.select_condition(o)
- f(o, s)
- elseif q.descent_condition(o)
- ds = descent_state(s)
- o, s = modify_stateful(f::F, (o, ds), q)
- o, merge_state(s, ds)
- else
- o, s
+@inline function modify_stateful(f::F, (obj, state), q::Q) where {F,Q<:Query}
+ let f=f, q=q
+ modify_stateful((obj, state), inner(q.optic)) do o, s
+ if (q::Q).select_condition(o)
+ (f::F)(o, s)
+ elseif (q::Q).descent_condition(o)
+ ds = descent_state(s)
+ o, ns = modify_stateful(f::F, (o, ds), q::Q)
+ o, merge_state(ds, ns)
+ else
+ o, s
+ end
  end
  end
 end
 
-maybesetproperties(state::GetAllState, obj, patch) = obj
+maybesetproperties(state::GetStates, obj, patch) = obj
 maybesetproperties(state::SetAllState, obj, patch) =
  maybesetproperties(state.change, state, obj, patch)
 maybesetproperties(::Changed, state::SetAllState, obj, patch) = setproperties(obj, patch)
@@ -516,8 +514,8 @@ anychanged(::Changed, ::Changed) = Changed()
 inner(optic) = optic
 inner(optic::ComposedOptic) = optic.inner
 
-outer(optic, o, state::GetAllState) = o
-outer(optic::ComposedOptic, o, state::GetAllState) = optic.outer(o)
+outer(optic, o, state::GetStates) = o
+outer(optic::ComposedOptic, o, state::GetStates) = optic.outer(o)
 outer(optic::ComposedOptic, o, state::SetAllState) = set(o, optic.outer, state.vals[state.itr])
 outer(optic, o, state::SetAllState) = state.vals[state.itr]
 
@@ -532,7 +530,7 @@ function (l::PropertyLens{field})(obj) where {field}
 end
 
 @inline function set(obj, l::PropertyLens{field}, val) where {field}
- patch = (;field => val)
+ patch = (; field => val)
  setproperties(obj, patch)
 end
 

test/test_queries.jl

@@ -1,46 +1,43 @@
 using Accessors, Test, BenchmarkTools, Static
 using Accessors: setall, getall, context
-
-obj = (7, (a=17.0, b=2.0f0), ("3", 4, 5.0), ((x=19, a=6.0,), [1,]))
+obj = (7, (a=17.0, b=2.0f0), ("3", 4, 5.0), ((x=19, a=6.0,)), [1])
 vals = (1.0, 2.0, 3.0, 4.0)
-
 # Fields is the default
 q = Query(; 
  select=x -> x isa NamedTuple, 
  descend=x -> x isa Tuple, 
  optic = (Accessors.@optic _.a) ∘ Accessors.Properties()
  # optic = Accessors.Properties()
 )
-
-println("getall")
 getall(obj, q)
+
 @code_native getall(obj, q)
 @code_warntype getall(obj, q)
 
 @benchmark getall($obj, $q)
 @test getall(obj, q) == (17.0, 6.0)
 
+# using ProfileView, Cthulhu
+# @descend getall(obj, q)
+# f(obj, q) = for i in 1:10000000 getall(obj, q) end
+# @profview f(obj, q)
+
 missings_obj = (a=missing, b=1, c=(d=missing, e=(f=missing, g=2)))
 @test getall(missings_obj, Query(ismissing)) === (missing, missing, missing)
 @benchmark getall($missings_obj, Query(ismissing))
 
-println("setall")
 # Need a wrapper so we don't have to pass in the starting iterator
 setall(obj, q, vals) 
 @benchmark setall($obj, $q, $vals)
+# using ProfileView
+# @profview for i in 1:1000000 setall(obj, q, vals) end
 @code_native setall(obj, q, vals)
 @code_warntype setall(obj, q, vals)
 
 # @btime Accessors.set($obj, $slowlens, $vals)
 @test setall(obj, q, vals) == 
- (7, (a=1.0, b=2.0f0), ("3", 4, 5.0), ((x=19, a=2.0,), [1]))
-
-using Cthulhu
-@descend getall(obj, q)
-# using ProfileView
-# @profview for i in 1:1000000 Accessors.set(obj, lens, vals) end
+ (7, (a=1.0, b=2.0f0), ("3", 4, 5.0), ((x=19, a=2.0,)), [1])
 
-println("unstable set")
 unstable_q = Accessors.Query(select=x -> x isa Float64 && x > 2, descend=x -> x isa NamedTuple)
 @btime setall($obj, $unstable_q, $vals)
 # slow_unstable_lens = Accessors.Query(; select=x -> x isa Number && x > 4, optic=Properties())
@@ -50,10 +47,15 @@ unstable_q = Accessors.Query(select=x -> x isa Float64 && x > 2, descend=x -> x
 @btime modify(x -> 10x, $obj, $q)
 
 # Context
-obj = (b=2, c=2)
-@test context((o, fn) -> fn, obj, q) == (:b, :c)
-@test context((o, fn) -> typeof(o), obj, q) == (typeof(obj), typeof(obj))
-@btime context((o, fn) -> fn, $obj, $q)
+q = Query(; 
+ select=x -> x isa Int, 
+ descend=x -> x isa NamedTuple, 
+ optic = Accessors.Properties()
+)
+obj2 = (1.0, :a, (b=2, c=2))
+@test context((o, fn) -> fn, obj2, q) == (:b, :c)
+@test context((o, fn) -> typeof(o), obj2, q) == (typeof(obj2[3]), typeof(obj2[3]))
+@btime context((o, fn) -> fn, $obj2, $q)
 
 # Macros
 @test (@getall (x for x in missings_obj if x isa Number)) == (1, 2)