A DenseArray describes a multidimensional array that has regular offets in memory. Implementing multidimensional arrays with the API flexibility of MATLAB (range indexing, colon indexing, etc.) is hard. Julia makes this easy by requiring only four methods to implement a DenseArray (five for copy semantics).

When an implementation satisfies the DenseArray interface, several behaviors are provided for free e.g. collect, find, maximum, sum, and many other functions in the standard library.

This post demonstrates what is necessary for a simple DenseArray implementation by constructing a lazy wrapper that squares elements of the underlying array when accessed. We call this lazy squaring wrapper LazySquares. Naturally, all mutations to the underlying array will be reflected in the lazy squaring wrapper.

Lazy Squares Specification

Given an existing multidimensional array xs, we would like to wrap it in a dense array using squares that will lazily square its contents. The desired behavior is illustrated in the example below.

julia> xs = [1, 2, 3]
3-element Array{Int64,1}:
 1
 2
 3

julia> xs_squared = square(xs)
3-element LazySquares.LazySquare{Int64,1}:
 1
 4
 9

Furthermore, changes to the underlying array xs should be reflected in the lazy squaring wrapper.

julia> xs[2] = 5
5

julia> xs_squared[2]
25

Lazy Squares Module and Type

We begin by creating a module. As per Julia convention, we name the module the plural form of the primary type provided by the module LazySquares since the primary type provided is LazySquare.

module LazySquares

type LazySquare{T <: Number, N} <: DenseArray{T, N}
    data::DenseArray{T, N}
end

end

It is clear from this interface that LazySquare will be a DenseArray with elements of type T while using N dimensions. The LazySquare type, as previously described, is a wrapper over an existing DenseArray; this makes it easy to extend by forwarding most methods to the underlying array.

Implementing Constructors and Methods

There are two types of methods to implement. First, we must implement constructors which are responsible for initializing the LazySquare type.

square{T <: Number, N}(xs::DenseArray{T, N})

Second, we implement methods for the corresponding functions required for DenseArray objects to operate normally:

Base.linearindexing(::Type{LazySquare})
Base.size(A::LazySquare)
Base.getindex(A,::LazySquare, i::Int)
Base.setindex!(A::LazySquare, v, i::Int)

The only non-obvious method to implement here is Base.linearindexing. This turns out to be very simple and we describe it in detail later.

Julia distinguishes between functions and methods. Functions are identified by their name. Functions have many methods which are definitions that are dispatched at runtime by the number of arguments passed to the function and the types of each argument. This is known as dynamic dispatch.

Square Constructor

Julia creates a default constructor with all fields as positional arguments with the appropriate types. We simply pass the array argument through to the default constructor of LazySquare.

square{T <: Number, N}(array::DenseArray{T, N}) = LazySquare(array)

Dense Array Methods

We have four methods to implement here.

We begin with the obscure Base.linearindexing. This method describes whether each location can be described with a single, scalar index efficiently. If this is the case, we assign the method the value of Base.LinearFast(); otherwise, we assign the method the value of Base.LinearSlow(). For DenseArray, it is obvious that we may use a single, scalar index due to its definition of having regular offsets in memory.

Base.linearindexing{T <: LazySquare}(::Type{T}) = Base.LinearFast()

Base.LinearSlow() is assumed by default. This implies that array requires indices for each dimension to locate an item in the array. This requires that implementing specialized getindex and setindex! methods for each of the number of dimensions supported.

Finally, we implement the standard array-like methods: size, getindex, and setindex!. The latter two methods describe array access and array assignment respectively.

"Return a tuple of sizes for each dimension"
Base.size{T <: Number, N}(array::LazySquare{T, N}) = size(array.data)

"Return an element of type `T` given an index `i`"
Base.getindex{T <: Number, N}(array::LazySquare{T, N}, i::Int) = array.data[i] * array.data[i]

"Assigns a value `v` of type `T` to the location identified by index `i`"
Base.setindex!{T <: Number, N}(array::LazySquare{T, N}, val::T, i::Int) = array.data[i] = val

It is important that we lazily square the contents of the underlying array in getindex as specified in our specification.

Testing

Given this implementation, it is also good practice to verify that our claims hold with unit tests. That is, we should verify that the elements of the underlying array are squared when accessed and that changes to the underlying array should be reflected in the lazy squaring wrapper.

using LazySquares
using Base.Test

xs = Int[1, 2, 3, 4]
xs_square = square(xs)

info("Verify that the elements are squared")
@test xs_square == xs .^ 2

info("Verify that getindex works")
@test xs_square[1] == xs[1] * xs[1]
@test xs_square[3] == xs[3] * xs[3]

info("Verify that setindex! works")
xs_square[2] = 5
@test xs_square[2] == 5 * 5
@test xs[2] == 5

Finally, we may verify that standard library functions for arrays also work as intended.

info("Verify that find works")
@test find(x -> x == 9, xs_square) == [3]

Further tests can be added for higher-dimensional arrays such as matrices and tensors.

Conclusion

Julia makes it easy to extend and implement multidimensional array-like behavior with only a concise set of methods. This level of generality enables the construction of SharedArray and DistributedArray types while reusing the same AbstractArray hierarchy.