Skip to content

Instantly share code, notes, and snippets.

@DNF2
Last active November 17, 2015 07:54
Show Gist options
  • Save DNF2/2d098eeefb6c00dd80fe to your computer and use it in GitHub Desktop.
Save DNF2/2d098eeefb6c00dd80fe to your computer and use it in GitHub Desktop.
A stream-like wrapper around iterable Julia objects. Based on https://github.com/andrewcooke/StatefulIterators.jl
module StatefulIterators_
export StatefulIterator, available, peek, reset!
import Base: eltype, start, next, done, read, copy
type StatefulIterator{T, S}
iter::T
state::S
end
StatefulIterator(itr) = StatefulIterator(itr, start(itr))
StatefulIterator(s::StatefulIterator) = StatefulIterator(s.iter, s.state)
copy(s::StatefulIterator) = StatefulIterator(s)
eltype{T,S}(::StatefulIterator{T,S}) = eltype(T)
eltype{T,S}(::Base.Take{StatefulIterator{T,S}}) = eltype(T)
start(s::StatefulIterator) = s
function next{T,S}(s::StatefulIterator{T,S}, ::StatefulIterator{T,S})
(item, s.state) = next(s.iter, s.state)
return item, s
end
done(s::StatefulIterator) = done(s.iter, s.state)
done{T,S}(s::StatefulIterator{T,S}, ::StatefulIterator{T,S}) = done(s)
read{T,S}(s::StatefulIterator{T,S}) = next(s, s)[1]
function read(s::StatefulIterator, dims::Int...)
a = Array(eltype(s), dims)
M = prod(dims)
if M == 0
return a
end
i = 0
for x in s
@inbounds a[i+=1] = x
if i == M
return a
end
end
return a[1:i] # Should possibly be an error. Asked for too many elements.
end
function peek(s::StatefulIterator, dims::Int...)
state = s.state
a = read(s, dims...)
s.state = state
return a
end
peek(s::StatefulIterator) = next(s.iter, s.state)[1]
@generated function available{T,S}(s::StatefulIterator{T,S})
# Returns the number of remaining bytes in stream.
if isbits(eltype(T)) && method_exists(length, Tuple{T})
# fast algorithm for iterators that support it
# NB! This fails for StepRange, since state can be > length!!
return :((length(s.iter) - s.state + 1) * sizeof(eltype(T)))
else
# fallback algorithm
return quote
state = s.state
siz = 0
for x in s
siz += sizeof(x)
end
s.state = state
return siz
end
end
end
function read{U}(s::StatefulIterator, ::Type{U})
N = cld(sizeof(U), sizeof(eltype(s)))
return reinterpret(U, read(s, N))[1]
end
function peek{U}(s::StatefulIterator, ::Type{U})
N = cld(sizeof(U), sizeof(eltype(s)))
return reinterpret(U, peek(s, N))[1]
end
function read{U}(s::StatefulIterator, ::Type{U}, dims...)
M = prod(dims)
N = cld(M * sizeof(U), sizeof(eltype(s)))
return reshape(reinterpret(U, read(s, N))[1:M], dims)
end
function peek{U}(s::StatefulIterator, ::Type{U}, dims...)
M = prod(dims)
N = cld(M * sizeof(U), sizeof(eltype(s)))
return reshape(reinterpret(U, peek(s, N))[1:M], dims)
end
reset!(s::StatefulIterator) = (s.state = start(s.iter))
end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment