Merge pull request #155 from torfjelde/tor/add-support-for-begin-indexing

Support for `begin` indexing

Author: jw3126

Project.toml

@@ -1,6 +1,6 @@
 name = "Setfield"
 uuid = "efcf1570-3423-57d1-acb7-fd33fddbac46"
-version = "0.7.0"
+version = "0.7.1"
 
 [deps]
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
@@ -26,4 +26,3 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
 test = ["Test", "Documenter", "PerformanceTestTools", "QuickTypes", "StaticArrays", "BenchmarkTools", "InteractiveUtils", "StaticNumbers"]
-

src/sugar.jl

@@ -61,26 +61,35 @@ foldtree(op, init, x) = op(init, x)
 foldtree(op, init, ex::Expr) =
     op(foldl((acc, x) -> foldtree(op, acc, x), ex.args; init=init), ex)
 
-need_dynamic_lens(ex) =
-    foldtree(false, ex) do yes, x
-        yes || x === :end || x === :_
-    end
+const HAS_BEGIN_INDEXING = VERSION ≥ v"1.5.0-DEV.666"
 
-replace_underscore(ex, to) = postwalk(x -> x === :_ ? to : x, ex)
+function need_dynamic_lens(ex)
+    return foldtree(false, ex) do yes, x
+        (yes || x === :end || (HAS_BEGIN_INDEXING && x === :begin) || x === :_)
+    end
+end
 
 function lower_index(collection::Symbol, index, dim)
     if isexpr(index, :call)
         return Expr(:call, lower_index.(collection, index.args, dim)...)
-    elseif index === :end
+    elseif (index === :end)
         if dim === nothing
             return :($(Base.lastindex)($collection))
         else
             return :($(Base.lastindex)($collection, $dim))
         end
+    elseif HAS_BEGIN_INDEXING && (index === :begin)
+        if dim === nothing
+            return :($(Base.firstindex)($collection))
+        else
+            return :($(Base.firstindex)($collection, $dim))
+        end
     end
     return index
 end
 
+replace_underscore(ex, to) = postwalk(x -> x === :_ ? to : x, ex)
+
 function parse_obj_lenses_composite(lensexprs::Vector)
     if isempty(lensexprs)
         return esc(:_), ()

test/dynamiclens_begin.jl

@@ -0,0 +1,18 @@
+l = @lens _[begin]
+@test l isa Setfield.DynamicIndexLens
+obj = (1,2,3)
+@test get(obj, l) == 1
+@test set(obj, l, true) == (true,2,3)
+
+l = @lens _[2*begin]
+@test l isa Setfield.DynamicIndexLens
+obj = (1,2,3)
+@test get(obj, l) == 2
+@test set(obj, l, true) == (1,true,3)
+
+one = 1
+plustwo(x) = x + 2
+l = @lens _.a[plustwo(begin) - one].b
+obj = (a=(1, (a=10, b=20), 3), b=4)
+@test get(obj, l) == 20
+@test set(obj, l, true) == (a=(1, (a=10, b=true), 3), b=4)

test/test_core.jl

@@ -281,6 +281,12 @@ end
     obj = (a=(1, (a=10, b=20), 3), b=4)
     @test get(obj, l) == 20
     @test set(obj, l, true) == (a=(1, (a=10, b=true), 3), b=4)
+
+    if Setfield.HAS_BEGIN_INDEXING
+        # Need to keep this in a separate file since `begin` won't parse
+        # on older Julia versions.
+        include("dynamiclens_begin.jl")
+    end
 end
 
 @testset "StaticNumbers" begin