Array.hs

-- | Core language array operations

module Feldspar.Core.Functions.Array
    ( arrayLen
    , getIx
    , setIx
    , getLength
    , RandomAccess (..)
    ) where



import Data.List

import Feldspar.Set
import Feldspar.Range
import Feldspar.Core.Types
import Feldspar.Core.Representation
import Feldspar.Core.Constructs
import Feldspar.Core.Functions.Num ()

import Feldspar.DSL.Expression
import Feldspar.DSL.Lambda
import Feldspar.DSL.Network

-- | Constructs an array of the given length and initialization.
arrayLen :: Type a => Data Length -> [a] -> Data [a]
arrayLen len = array (dataSize len :> universal)
  -- TODO This function is a temporary solution.

evalGetIx :: Range Length -> [a] -> Index -> a
evalGetIx l as i
    | not (i `inRange` r) = error "getIx: index out of bounds"
    | i >= la             = error "getIx: reading garbage"
    | otherwise           = genericIndex as i
  where
    la = genericLength as
    r  = rangeByRange 0 (l-1)

-- | Look up an index in an array (see also '!')
getIx :: Type a => Data [a] -> Data Index -> Data a
getIx arr = function2 "(!)" sizeProp (evalGetIx l) arr
  where
    sizeProp (_:>aSize) _ = aSize
    l:>_ = dataSize arr

-- | Array update
setIx
    :: Type a
    => Data [a]    -- ^ Source array
    -> Data Index  -- ^ Index to replace
    -> Data a      -- ^ New value
    -> Data [a]
setIx arr i a =
        nodeData (dataSize arr)
     $  Inject (Node SetIx)
    :$: toEdge i
    :$: toEdge a
    :$: toEdge arr

-- | Array length
getLength :: Type a => Data [a] -> Data Length
getLength arr = case undoEdge (unData arr) of
    Inject (Node Parallel) :$: len :$: _ :$: arr' -> Data len + getLength (Data arr')
    Inject (Node SetLength) :$: len :$: _ -> Data len
    Inject (Node SetIx) :$: _ :$: _ :$: arr' -> getLength (Data arr')
    _ -> case dataSize arr of
        (Range l b :> _)
            | l == b    -> value l
        otherwise       -> function1 "length" sizeProp genericLength arr
  where
    sizeProp (lSize:>_) = lSize

infixl 9 !

class RandomAccess a
  where
    -- | The type of elements in a random access structure
    type Element a

    -- | Index lookup in a random access structure
    (!) :: a -> Data Index -> Element a

instance Type a => RandomAccess (Data [a])
  where
    type Element (Data [a]) = Data a
    (!) = getIx