refactor controldo - fix a corner case (#41)

* update * fix tests * fix tests
QuantumBFS · Jun 22, 2022 · a1e07e6 · a1e07e6
1 parent 0a009a4
commit a1e07e6
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/src/iterate_control.jl b/src/iterate_control.jl
@@ -110,21 +110,19 @@ function group_shift!(nbits::Int, positions::AbstractVector{Int})
     factors = Int[]
     k_prv = 0
     i = 0
-    m = length(positions)
-    # 2 456 9 ...
-    for (l,k) in enumerate(positions)
+    for k in positions
         if k != k_prv + 1
             push!(factors, 1<<(k_prv-i))
             gap = k - k_prv-1
             push!(masks, bmask(i+1:i+gap))
             i += gap
         end
         k_prv = k
-        # the last block
-        if l == m && m != nbits
-            push!(factors, 1<<(k-i))
-            push!(masks, bmask(i+1:nbits))
-        end
+    end
+    # the last block
+    if i != nbits
+        push!(factors, 1<<(k_prv-i))
+        push!(masks, bmask(i+1:nbits))
     end
     return masks, factors
 end
diff --git a/test/iterate_control.jl b/test/iterate_control.jl
@@ -21,6 +21,10 @@ mulrow!(v::AbstractVector, i::Int, f) = (v[i] *= f; v)
     @test group_shift!(5, [1, 2, 5]) == ([3, 28], [4, 8])
     @test group_shift!(5, [2, 3]) == ([1, 30], [1, 4])
     @test group_shift!(5, [1, 3, 5]) == ([1, 2, 28], [2, 4, 8])
+    @test group_shift!(3, [1, 2, 3]) == ([7], [8])
+    @test group_shift!(3, Int[]) == (Int[7], Int[1])
+    @test [itercontrol(3, [1, 2, 3], [0,1,0])...] == [2]
+    @test [itercontrol(3, Int[], Int[])...] == [0,1,2,3,4,5,6,7]
 end
 
 @testset "test IterControl" begin