From 6b6964cae4044debd8f3826efa967f2a758faede Mon Sep 17 00:00:00 2001
From: Jinguo Liu <cacate0129@gmail.com>
Date: Fri, 1 Jul 2022 16:01:28 -0400
Subject: [PATCH] update (#43)

---
 .gitignore             |  1 +
 src/bit_operations.jl  | 13 ++++++++-----
 test/bit_operations.jl |  6 ++++++
 3 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/.gitignore b/.gitignore
index 93a6cb9..0c0c08e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -9,3 +9,4 @@ docs/Manifest.toml
 docs/build
 docs/src/assets/indigo.css
 .vscode
+*.swp
diff --git a/src/bit_operations.jl b/src/bit_operations.jl
index 6162d93..eef1e4b 100644
--- a/src/bit_operations.jl
+++ b/src/bit_operations.jl
@@ -307,14 +307,17 @@ end
 const IntIterator{T} = Union{NTuple{<:Any,T},Vector{T},T,UnitRange{T}} where {T<:Integer}
 
 """
-    controller(cbits, cvals) -> Function
+    controller([T=Int, ]cbits, cvals) -> Function
 
 Return a function that checks whether a basis at `cbits` takes specific value `cvals`.
 """
-function controller(cbits::IntIterator{Int}, cvals::IntIterator{Int})
-    do_mask = bmask(cbits)
+function controller(cbits::IntIterator, cvals::IntIterator)
+    controller(Int, cbits, cvals)
+end
+function controller(::Type{T}, cbits::IntIterator, cvals::IntIterator) where T
+    do_mask = bmask(T, cbits...)
     target =
-        length(cvals) == 0 ? 0 :
-        mapreduce(xy -> (xy[2] == 1 ? 1 << (xy[1] - 1) : 0), |, zip(cbits, cvals))
+        length(cvals) == 0 ? zero(T) :
+        mapreduce(xy -> (xy[2] == 1 ? one(T) << T(xy[1] - 1) : zero(T)), |, zip(cbits, cvals))
     return b -> ismatch(b, do_mask, target)
 end
diff --git a/test/bit_operations.jl b/test/bit_operations.jl
index 672b20c..92ec641 100644
--- a/test/bit_operations.jl
+++ b/test/bit_operations.jl
@@ -60,3 +60,9 @@ end
 @testset "random" begin
     @test rand(BitStr64{5}) isa BitStr64{5}
 end
+
+@testset "controller" begin
+    c = controller(Int128, [2,3], [0, 1])
+    @test c(Int128(3)) == false
+    @test c(Int128(0b101)) == true
+end