Paste: multi-arrays

! (c)2010 Kobi Lurie, Joe Groff bsd license
USING: accessors kernel math math.order sequences ;
FROM: sequences.private => nth-unsafe set-nth-unsafe ;
IN: sequences.multi

MIXIN: multi-sequence

GENERIC: multi-dim ( multi-array -- dim )

: multi-bounds-check? ( index multi-array -- ? )
    multi-dim [ 0 swap between? ] 2all? ; inline

: multi-bounds-check ( index multi-array -- index multi-array )
    2dup multi-bounds-check? [ bounds-error ] unless ; inline

GENERIC: multi-nth ( index multi-array -- elt )
GENERIC: multi-set-nth ( elt index multi-array -- )

<PRIVATE
GENERIC: multi-nth-unsafe ( index multi-array -- elt )
GENERIC: multi-set-nth-unsafe ( elt index multi-array -- )

M: multi-sequence multi-nth-unsafe multi-nth ; inline
M: multi-sequence multi-set-nth-unsafe multi-set-nth ; inline
PRIVATE>

M: multi-sequence multi-nth ( index multi-array -- elt )
    multi-bounds-check multi-nth-unsafe ; inline
M: multi-sequence multi-set-nth ( index multi-array -- elt )
    multi-bounds-check multi-set-nth-unsafe ; inline

TUPLE: multi-array
    { origin read-only }
    { dim read-only }
    { strides read-only }
    { underlying read-only } ;
INSTANCE: multi-array multi-sequence

: row-major ( dim -- strides )
    1 [ * ] accumulate nip ; inline

: col-major ( dim -- strides )
    row-major reverse ; inline

: (multi-array-like) ( origin dim strides exemplar -- multi-array )
    [ over product ] dip new-sequence
    multi-array boa ; inline

: <multi-array-like> ( dim exemplar -- multi-array ) 
    [ 0 swap dup row-major ] dip (multi-array-like) ; inline

: <multi-array> ( dim -- multi-array )
    { } <multi-array-like> ; inline

: multi@ ( index strides -- n )
    [ * ] [ + ] 2map-reduce ; inline

M: multi-array multi-dim dim>> ; inline
M: multi-array multi-nth-unsafe [ strides>> multi@ ] [ underlying>> ] bi nth-unsafe ; inline
M: multi-array multi-set-nth-unsafe [ strides>> multi@ ] [ underlying>> ] bi set-nth-unsafe ; inline