module Array where
import Data.Word

import LLVM.Core
--import LLVM.ExecutionEngine
import LLVM.Util.Loop
import LLVM.Util.Optimize

cg :: CodeGenModule (Function (Double -> IO (Ptr Double)))
cg = do
    dotProd <- createFunction InternalLinkage $ \ size aPtr aStride bPtr bStride -> do
        r <- forLoop (valueOf 0) size (valueOf 0) $ \ i s -> do
	    ai <- mul aStride i
	    bi <- mul bStride i
            ap <- getElementPtr aPtr (ai, ())
            bp <- getElementPtr bPtr (bi, ())
            a <- load ap
            b <- load bp
            ab <- mul a b
            add (s :: Value Double) ab
	ret r
    let _ = dotProd :: Function (Word32 -> Ptr Double -> Word32 -> Ptr Double -> Word32 -> IO Double)

    -- multiply a:[n x m], b:[m x l]
    matMul <- createFunction InternalLinkage $ \ n m l aPtr bPtr cPtr -> do
        forLoop (valueOf 0) n () $ \ ni () -> do
           forLoop (valueOf 0) l () $ \ li () -> do
	      ni' <- mul ni m
	      row <- getElementPtr aPtr (ni', ())
	      col <- getElementPtr bPtr (li, ())
              x <- call dotProd m row (valueOf 1) col m
	      j <- add ni' li
	      p <- getElementPtr cPtr (j, ())
	      store x p
	      return ()
        ret ()
    let _ = matMul :: Function (Word32 -> Word32 -> Word32 -> Ptr Double -> Ptr Double -> Ptr Double -> IO ())

    let fillArray _ [] = return ()
        fillArray ptr (x:xs) = do store x ptr; ptr' <- getElementPtr ptr (1::Word32,()); fillArray ptr' xs

    test <- createNamedFunction ExternalLinkage "test" $ \ x -> do
        a <- arrayMalloc (4 :: Word32)
	fillArray a $ map valueOf [1,2,3,4]
	b <- arrayMalloc (4 :: Word32)
	fillArray b [x,x,x,x]
	c <- arrayMalloc (4 :: Word32)
	_ <- call matMul (valueOf 2) (valueOf 2) (valueOf 2) a b c
	ret c
    let _ = test :: Function (Double -> IO (Ptr Double))

    return test

main :: IO ()
main = do
    -- Initialize jitter
    initializeNativeTarget
    m <- newModule
    _f <- defineModule m cg
    writeBitcodeToFile "Arr.bc" m
    _ <- optimizeModule 3 m
    writeBitcodeToFile "Arr-opt.bc" m