From ca29f11d46d90d36a157cf8c1261229b8350b2ec Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jinguo=20Liu=20=28=E5=88=98=E9=87=91=E5=9B=BD=29?=
 <cacate0129@gmail.com>
Date: Sat, 25 May 2024 01:30:36 +0800
Subject: [PATCH] Improve readbit (#51)

* update

* fix tests
---
 src/DitStr.jl        |  2 ++
 src/longlonguint.jl  |  6 +++++-
 test/longlonguint.jl | 15 +++++++++++++++
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/src/DitStr.jl b/src/DitStr.jl
index a7ba44a..669e19f 100644
--- a/src/DitStr.jl
+++ b/src/DitStr.jl
@@ -233,6 +233,8 @@ One can specify the value of the nonzero entry by inputing a pair.
 onehot(::Type{T}, n::DitStr{D,N,T1}; nbatch=nothing) where {D,T, N,T1} = _onehot(T, D^N, buffer(n)+1; nbatch)
 onehot(n::DitStr; nbatch=nothing) = onehot(ComplexF64, n; nbatch)
 
+readbit(x::DitStr{D, N, LongLongUInt{C}}, loc::Int) where {D, N, C} = readbit(x.buf, loc)
+
 ########## @dit_str macro ##############
 """
     @dit_str -> DitStr64
diff --git a/src/longlonguint.jl b/src/longlonguint.jl
index ca11cf1..709cc9c 100644
--- a/src/longlonguint.jl
+++ b/src/longlonguint.jl
@@ -22,7 +22,7 @@ Base.UInt(x::LongLongUInt{1}) = x.content[1]
 Base.zero(::Type{LongLongUInt{C}}) where {C} = LongLongUInt{C}(ntuple(_->UInt(0), Val{C}()))
 Base.zero(::LongLongUInt{C}) where {C} = zero(LongLongUInt{C})
 # convert from integers
-LongLongUInt{C}(x::T) where {C, T<:Integer} = LongLongUInt{C}(ntuple(i->i==1 ? UInt(x) : zero(UInt), Val{C}()))
+LongLongUInt{C}(x::T) where {C, T<:Integer} = LongLongUInt{C}(ntuple(i->i==C ? UInt(x) : zero(UInt), Val{C}()))
 Base.promote_type(::Type{LongLongUInt{C}}, ::Type{Int}) where {C} = LongLongUInt{C}
 Base.promote_type(::Type{LongLongUInt{C}}, ::Type{UInt}) where {C} = LongLongUInt{C}
 function Base.mod(x::LongLongUInt{C}, D::Int) where {C}
@@ -51,6 +51,10 @@ function Base.:(<<)(x::LongLongUInt{C}, y::Int) where C
         end
     )
 end
+function readbit(x::LongLongUInt{C}, loc::Int) where {C}
+    k = (loc-1) ÷ bsizeof(UInt)
+    return readbit(x.content[C-k], loc - k*bsizeof(UInt))
+end
 function indicator(::Type{LongLongUInt{C}}, i::Int) where C
     k = (i-1) ÷ bsizeof(UInt)
     LongLongUInt{C}(ntuple(j->j==C-k ? indicator(UInt, i-k*bsizeof(UInt)) : zero(UInt), Val{C}()))
diff --git a/test/longlonguint.jl b/test/longlonguint.jl
index 1d9ee6a..212ae40 100644
--- a/test/longlonguint.jl
+++ b/test/longlonguint.jl
@@ -14,6 +14,7 @@ using Test, BitBasis
     @test ~x == LongLongUInt((~UInt(3), ~UInt(6)))
 
     y = LongLongUInt((5, 7))
+    @test one(y) == LongLongUInt((0, 1))
     @test x & y == LongLongUInt((1, 6))
     @test x | y == LongLongUInt((7, 7))
     @test x ⊻ y == LongLongUInt((6, 1))
@@ -62,4 +63,18 @@ end
     @test indicator(LongLongUInt{1}, 64) == LongLongUInt((UInt(1)<<63,))
     @test indicator(LongLongUInt{2}, 64) == LongLongUInt((zero(UInt), UInt(1)<<63))
     @test indicator(LongLongUInt{2}, 65) == LongLongUInt((1, 0))
+end
+
+@testset "takebit" begin
+    x = LongLongUInt((3, 6))
+    @test readbit(x, 1) == 0
+    @test readbit(x, 2) == 1
+    @test readbit(x, 3) == 1
+    @test readbit(x, 4) == 0
+
+    @test readbit(x, 64) == 0
+    @test readbit(x, 65) == 1
+    @test readbit(x, 66) == 1
+    @test readbit(x, 67) == 0
+    @test readbit(BitStr{78}(x), 66) == 1
 end
\ No newline at end of file