-
Notifications
You must be signed in to change notification settings - Fork 2
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
91a30eb
commit ba93856
Showing
10 changed files
with
943 additions
and
189 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,14 @@ | ||
module UnrolledUtilitiesStaticArraysExt | ||
|
||
import UnrolledUtilities | ||
import StaticArrays: SVector, MVector | ||
|
||
UnrolledUtilities.type_length(::Type{<:SVector{N}}) where {N} = N | ||
UnrolledUtilities.output_type_for_promotion(::SVector) = SVector | ||
UnrolledUtilities.constructor_from_tuple(::Type{SVector}) = SVector | ||
|
||
UnrolledUtilities.type_length(::Type{<:MVector{N}}) where {N} = N | ||
UnrolledUtilities.output_type_for_promotion(::MVector) = MVector | ||
UnrolledUtilities.constructor_from_tuple(::Type{MVector}) = MVector | ||
|
||
end |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,144 @@ | ||
""" | ||
BitSequence{N, [U]}(f) | ||
BitSequence{N, [U]}([bit]) | ||
A statically-sized analogue of `BitVector` with `Unsigned` chunks of type `U`, | ||
which can be constructed using either a function `f(n)` or a constant `bit`. By | ||
default, `U` is set to `UInt8` and `bit` is set to `false`. | ||
This iterator can only store `Bool`s, so its `output_type_for_promotion` is a | ||
`ConditionalOutputType`. Efficient methods are provided for `unrolled_map`, | ||
`unrolled_accumulate`, `unrolled_take`, and `unrolled_drop`, though the methods | ||
for `unrolled_map` and `unrolled_accumulate` only apply when the first items in | ||
their outputs are `Bool`s. All other unrolled functions that need to generate | ||
iterators use output types that are not `BitSequence`s. | ||
""" | ||
struct BitSequence{N, U <: Unsigned, I <: NTuple{<:Any, U}} <: StaticSequence{N} | ||
ints::I | ||
end | ||
BitSequence{N, U}(ints) where {N, U} = BitSequence{N, U, typeof(ints)}(ints) | ||
BitSequence{N}(args...) where {N} = BitSequence{N, UInt8}(args...) | ||
|
||
function BitSequence{N, U}(bit::Bool = false) where {N, U} | ||
n_bits_per_int = 8 * sizeof(U) | ||
n_ints = cld(N, n_bits_per_int) | ||
int = bit ? ~zero(U) : zero(U) | ||
ints = ntuple(_ -> int, Val(n_ints)) | ||
return BitSequence{N, U}(ints) | ||
end | ||
|
||
function BitSequence{N, U}(f::Function) where {N, U} | ||
n_bits_per_int = 8 * sizeof(U) | ||
n_ints = cld(N, n_bits_per_int) | ||
ints = ntuple(Val(n_ints)) do int_index | ||
@inline | ||
first_index = n_bits_per_int * (int_index - 1) + 1 | ||
unrolled_reduce( | ||
LazySequence{min(n_bits_per_int, N - first_index + 1)}(0); | ||
init = zero(U), | ||
) do int, bit_offset | ||
int | U(f(first_index + bit_offset)::Bool) << bit_offset | ||
end | ||
end | ||
return BitSequence{N, U}(ints) | ||
end | ||
|
||
@inline function int_index_and_bit_offset(::Type{U}, n) where {U} | ||
int_offset, bit_offset = divrem(n - 1, 8 * sizeof(U)) | ||
return (int_offset + 1, bit_offset) | ||
end | ||
|
||
@inline function Base.getindex(itr::BitSequence{<:Any, U}, n::Integer) where {U} | ||
int_index, bit_offset = int_index_and_bit_offset(U, n) | ||
int = itr.ints[int_index] | ||
return Bool(int >> bit_offset & one(int)) | ||
end | ||
|
||
@inline function Base.setindex( | ||
itr::BitSequence{N, U}, | ||
bit::Bool, | ||
n::Integer, | ||
) where {N, U} | ||
int_index, bit_offset = int_index_and_bit_offset(U, n) | ||
int = itr.ints[int_index] | ||
int′ = int & ~(one(int) << bit_offset) | U(bit) << bit_offset | ||
return BitSequence{N, U}(Base.setindex(itr.ints, int′, int_index)) | ||
end | ||
|
||
output_type_for_promotion(::BitSequence{<:Any, U}) where {U} = | ||
ConditionalOutputType(Bool, BitSequence{<:Any, U}) | ||
|
||
@inline function unrolled_map_into( | ||
::Type{BitSequence{<:Any, U}}, | ||
f, | ||
itrs..., | ||
) where {U} | ||
lazy_itr = Iterators.map(f, itrs...) | ||
N = inferred_length(lazy_itr) | ||
return BitSequence{N, U}(Base.Fix1(extended_getindex, lazy_itr)) | ||
end | ||
|
||
@inline function unrolled_accumulate_into( | ||
::Type{BitSequence{<:Any, U}}, | ||
op, | ||
itr, | ||
init, | ||
transform, | ||
) where {U} | ||
N = inferred_length(itr) | ||
(N == 0 && init isa NoInit) && | ||
error("unrolled_accumulate requires an init value for empty iterators") | ||
n_bits_per_int = 8 * sizeof(U) | ||
n_ints = cld(N, n_bits_per_int) | ||
ints = unrolled_accumulate_into_tuple( | ||
LazySequence{n_ints}(); | ||
init = (nothing, init), | ||
transform = first, | ||
) do (_, init_value_for_new_int), int_index | ||
@inline | ||
first_index = n_bits_per_int * (int_index - 1) + 1 | ||
unrolled_reduce( | ||
LazySequence{min(n_bits_per_int, N - first_index + 1)}(0); | ||
init = (zero(U), init_value_for_new_int), | ||
) do (int, prev_value), bit_offset | ||
item = extended_getindex(itr, first_index + bit_offset) | ||
new_value = | ||
first_index + bit_offset == 1 && prev_value isa NoInit ? | ||
item : op(prev_value, item) | ||
(int | U(transform(new_value)::Bool) << bit_offset, new_value) | ||
end | ||
end | ||
return BitSequence{N, U}(ints) | ||
end | ||
|
||
@inline function unrolled_take( | ||
itr::BitSequence{<:Any, U}, | ||
::Val{N}, | ||
) where {N, U} | ||
n_bits_per_int = 8 * sizeof(U) | ||
n_ints = cld(N, n_bits_per_int) | ||
ints = unrolled_take(itr.ints, Val(n_ints)) | ||
return BitSequence{N, U}(ints) | ||
end | ||
|
||
@inline function unrolled_drop( | ||
itr::BitSequence{N_old, U}, | ||
::Val{N}, | ||
) where {N_old, N, U} | ||
n_bits_per_int = 8 * sizeof(U) | ||
n_ints = cld(N_old - N, n_bits_per_int) | ||
n_dropped_ints = length(itr.ints) - n_ints | ||
bit_offset = N - n_bits_per_int * n_dropped_ints | ||
ints_without_offset = unrolled_drop(itr.ints, Val(n_dropped_ints)) | ||
ints = if bit_offset == 0 | ||
ints_without_offset | ||
else | ||
cur_ints = ints_without_offset | ||
next_ints = unrolled_push(unrolled_drop(cur_ints, Val(1)), nothing) | ||
unrolled_map_into_tuple(cur_ints, next_ints) do cur_int, next_int | ||
isnothing(next_int) ? cur_int >> bit_offset : | ||
cur_int >> bit_offset | next_int << (n_bits_per_int - bit_offset) | ||
end | ||
end | ||
return BitSequence{N_old - N, U}(ints) | ||
end |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,50 @@ | ||
""" | ||
LazySequence{N}(f) | ||
LazySequence{N}([start]) | ||
A lazy analogue of `ntuple(f, Val(N))`, or a lazy and statically-sized analogue | ||
of `start:(start - 1 + N)`. By default, `start` is set to 1. | ||
Efficient methods are provided for `unrolled_take` and `unrolled_drop`. All | ||
other unrolled functions that need to generate iterators use output types that | ||
are not `LazySequence`s. | ||
""" | ||
struct LazySequence{N, F} <: StaticSequence{N} | ||
f::F | ||
end | ||
LazySequence{N}(f::Function = identity) where {N} = | ||
LazySequence{N, typeof(f)}(f) | ||
LazySequence{N}(start::Number) where {N} = | ||
LazySequence{N}(Base.Fix1(+, start - one(start))) | ||
|
||
@inline Base.getindex(itr::LazySequence, n::Integer) = itr.f(n) | ||
|
||
output_type_for_promotion(::LazySequence) = NoOutputType() | ||
|
||
output_promote_rule(::Type{LazySequence}, ::Type{O}) where {O} = O | ||
|
||
@inline unrolled_take(itr::LazySequence, ::Val{N}) where {N} = | ||
LazySequence{N}(itr.f) | ||
|
||
@inline unrolled_drop(itr::LazySequence{N_old}, ::Val{N}) where {N_old, N} = | ||
LazySequence{N_old - N}(n -> itr.f(n + N)) | ||
|
||
""" | ||
StaticOneTo(N) | ||
A lazy analogue of `Base.OneTo(N)`. | ||
An efficient method is provided for `unrolled_take`. All other unrolled | ||
functions that need to generate iterators use output types that are not | ||
`StaticOneTo`s. | ||
""" | ||
struct StaticOneTo{N} <: StaticSequence{N} end | ||
StaticOneTo(N) = StaticOneTo{N}() | ||
|
||
@inline Base.getindex(::StaticOneTo, n::Integer) = n | ||
|
||
output_type_for_promotion(::StaticOneTo) = NoOutputType() | ||
|
||
output_promote_rule(::Type{StaticOneTo}, ::Type{O}) where {O} = O | ||
|
||
@inline unrolled_take(::StaticOneTo, ::Val{N}) where {N} = StaticOneTo(N) |
Oops, something went wrong.