module LLVM.Extra.Array (
   size,
   assemble,
   extractAll,
   map,
   ) where

import qualified LLVM.Extra.Class as Class

import qualified LLVM.Core as LLVM
import LLVM.Core (Value, Array, CodeGenFunction, )

import qualified Data.TypeLevel.Num as TypeNum
import Control.Monad.HT ((<=<), )
import Control.Monad (foldM, )
import qualified Data.List as List

import Data.Word (Word32, )

import Prelude hiding
          (Real, truncate, floor, round,
           map, zipWith, iterate, replicate, reverse, concat, sum, )


-- * target independent functions

size ::
   (TypeNum.Nat n) =>
   Value (Array n a) -> Int
size =
   let sz :: (TypeNum.Nat n) => n -> Value (Array n a) -> Int
       sz n _ = TypeNum.toInt n
   in  sz undefined

{- |
construct an array out of single elements

You must assert that the length of the list matches the array size.

This can be considered the inverse of 'extractAll'.
-}
assemble ::
   (TypeNum.Nat n, LLVM.IsFirstClass a, LLVM.IsSized a s) =>
   [Value a] -> CodeGenFunction r (Value (Array n a))
assemble =
   foldM (\v (k,x) -> LLVM.insertvalue v x (k::Word32)) Class.undefTuple .
   List.zip [0..]

{- |
provide the elements of an array as a list of individual virtual registers

This can be considered the inverse of 'assemble'.
-}
extractAll ::
   (TypeNum.Nat n, LLVM.IsFirstClass a, LLVM.IsSized a s) =>
   Value (Array n a) -> LLVM.CodeGenFunction r [Value a]
extractAll x =
   mapM
      (LLVM.extractvalue x)
      (take (size x) [(0::Word32)..])

{- |
The loop is unrolled,
since 'LLVM.insertvalue' and 'LLVM.extractvalue' expect constant indices.
-}
map ::
   (TypeNum.Nat n,
    LLVM.IsFirstClass a, LLVM.IsSized a asize,
    LLVM.IsFirstClass b, LLVM.IsSized b bsize) =>
   (Value a -> CodeGenFunction r (Value b)) ->
   (Value (Array n a) -> CodeGenFunction r (Value (Array n b)))
map f =
   assemble <=< mapM f <=< extractAll