Make axisfor inferrable

tkf · Jul 19, 2019 · 28d7fb4 · 28d7fb4
1 parent df6611b
commit 28d7fb4
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 5 deletions.
diff --git a/src/reduce.jl b/src/reduce.jl
@@ -135,11 +135,11 @@ end
 
 axisfor(patterns, i) =
     foldlargs(nothing, patterns...) do _, p::AccessPattern
-        foldlargs(nothing, ntuple(identity, length(p.indices))...) do _, n
-            if p.indices[n] === i
+        foldlargs(1, p.indices...) do n, j
+            if j === i
                 reduced(axes(p.indexable)[n])
             else
-                nothing
+                n + 1
             end
         end
     end |> unreduced

diff --git a/test/preamble.jl b/test/preamble.jl
@@ -2,7 +2,7 @@ using Base.Broadcast: broadcasted, instantiate
 using Test
 using MacroTools
 using NDReducibles
-using NDReducibles: Index, plan
+using NDReducibles: AccessPattern, Index, plan, axisfor
 
 _indices(idxs::String) = Index.(Tuple(Symbol.(collect(idxs))))
 

diff --git a/test/test_inference.jl b/test/test_inference.jl
@@ -12,7 +12,7 @@ f2() = plan(
     ND(2) => (:i, :j),
 )
 
-@testset begin
+@testset "plan" begin
     if VERSION < v"1.2-"
         @test_broken_inferred f1()
         @test_broken_inferred f2()
@@ -22,4 +22,13 @@ f2() = plan(
     end
 end
 
+@testset "axisfor" begin
+    patterns = AccessPattern.((
+        ND(1) => (:i,),
+        ND(2) => (:i, :j),
+    ))
+    @test_inferred axisfor(patterns, Index(:i))
+    @test_inferred axisfor(patterns, Index(:j))
+end
+
 end  # module