<commit_before>module Iter

import Base.assign, Base.done, Base.next, Base.ref, Base.start

parent(A::Array) = A
parent(A::SubArray) = A.parent
first_index(A::Array) = 1
first_index(A::SubArray) = A.first_index

function Iterator(A::StridedArray, dims, slices)
    notdims = setdiff(1:ndims(A), [dims...])
    if length(slices) != length(notdims)
        error("The dimensionalities do not add up")
    end
    sz = size(A)
    s = strides(A)
    N = length(dims)
    f = first_index(A)
    for i = 1:length(slices)
        if ~isa(slices[i], Integer)
            error("slices must be integers")
        end
        if !(1 <= slices[i] <= sz[i])
            throw(BoundsError())
        end
        f += (slices[i]-1)*s[notdims[i]]
    end
    if N == 2
        return Iterator2D(sz[dims[1]], sz[dims[2]], s[dims[1]], s[dims[2]], f)
    else
        error(N, " dimensions not yet supported")
    end
end
Iterator(A::StridedArray) = Iterator(A, 1:ndims(A), ())

#### Iterator definitions ####
# Iterator stores size information, IteratorState holds the current state

abstract IteratorState

# 2D
immutable Iterator2D
    ni::Int
    nj::Int
    stridei::Int
    stridej::Int
    first_index::Int
end

immutable Iterator2DState <: IteratorState
    i::Int
    j::Int
    index::Int
end

start(iter::Iterator2D) = Iterator2DState(1, 1, iter.first_index)

next(iter::Iterator2D, state::Iterator2DState) = (state.i+1 <= iter.ni) ? (state, Iterator2DState(state.i+1, state.j, state.index + iter.stridei)) : (state, Iterator2DState(1, state.j+1, iter.first_index+state.j*iter.stridej))

done(iter::Iterator2D, state::Iterator2DState) = state.j > iter.nj


#### ref/assign ####

ref(A::StridedArray, state::IteratorState) = parent(A)[state.index]

assign{T}(A::StridedArray{T}, x::T, state::IteratorState) = assign(parent(A), x, state.index)

end<commit_msg>Add interface for 1D iteration<commit_after>module Iter

import Base.assign, Base.done, Base.next, Base.ref, Base.start

export Iterator

parent(A::Array) = A
parent(A::SubArray) = A.parent
first_index(A::Array) = 1
first_index(A::SubArray) = A.first_index

function Iterator(A::StridedArray, dims, slices)
    notdims = setdiff(1:ndims(A), [dims...])
    if length(slices) != length(notdims)
        error("The dimensionalities do not add up")
    end
    sz = size(A)
    s = strides(A)
    N = length(dims)
    f = first_index(A)
    for i = 1:length(slices)
        if ~isa(slices[i], Integer)
            error("slices must be integers")
        end
        if !(1 <= slices[i] <= sz[i])
            throw(BoundsError())
        end
        f += (slices[i]-1)*s[notdims[i]]
    end
    if N == 1
        return Iterator1D(sz[dims[1]], s[dims[1]], f)
    elseif N == 2
        return Iterator2D(sz[dims[1]], sz[dims[2]], s[dims[1]], s[dims[2]], f)
    else
        error(N, " dimensions not yet supported")
    end
end
Iterator(A::StridedArray) = Iterator(A, 1:ndims(A), ())

#### Iterator definitions ####
# Iterator stores size information, IteratorState holds the current state

abstract IteratorState

# 1D
immutable Iterator1D
    ni::Int
    stridei::Int
    first_index::Int
end

immutable Iterator1DState <: IteratorState
    i::Int
    index::Int
end

start(iter::Iterator1D) = Iterator1DState(1, iter.first_index)

# next(iter::Iterator1D, state::Iterator1DState) = state, Iterator1DState(state.i+1, state.index + iter.stridei)
next(iter::Iterator1D, state::Iterator1DState) = state, Iterator1DState(state.i+1, state.index+1)

done(iter::Iterator1D, state::Iterator1DState) = state.i > iter.ni



# 2D
immutable Iterator2D
    ni::Int
    nj::Int
    stridei::Int
    stridej::Int
    first_index::Int
end

immutable Iterator2DState <: IteratorState
    i::Int
    j::Int
    index::Int
end

start(iter::Iterator2D) = Iterator2DState(1, 1, iter.first_index)

next(iter::Iterator2D, state::Iterator2DState) = (state.i+1 <= iter.ni) ? (state, Iterator2DState(state.i+1, state.j, state.index + iter.stridei)) : (state, Iterator2DState(1, state.j+1, iter.first_index+state.j*iter.stridej))

done(iter::Iterator2D, state::Iterator2DState) = state.j > iter.nj


#### ref/assign ####

ref(A::StridedArray, state::IteratorState) = parent(A)[state.index]

assign{T}(A::StridedArray{T}, x::T, state::IteratorState) = assign(parent(A), x, state.index)

end
{'fraction_non_alphanumeric': 0.15611227093481791, 'fraction_numerical': 0.01855717930874507, 'mean_word_length': 3.928, 'pattern_counts': {'":': 0, '<': 12, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4775271', 'n_tokens_mistral': 1595, 'n_tokens_neox': 1596, 'n_words': 327}