-- Generating imperative code with loops and arrays
-- This code is a safer version of the loop-tiling
-- matrix-vector--multiplication example of talk-problems.ml. 
-- See the latter file for the motivation and explanation.

module TSLoop where

import TSCore
import TSCPST

import Control.Applicative
import Data.Array.MArray
import Data.Array.IO


-- A sample matrix
sample_a :: IO (IOUArray (Int,Int) Int)
sample_a = do
  a <- newArray ((0,0),(4,9)) 0
  sequence [writeArray a (i,j) (i+j) |
	    i <- [0..4], j <- [0..9]]
  return a

test_a = sample_a >>= getAssocs >>= print

-- A sample vector
sample_v :: IO (IOUArray Int Int)
sample_v = do
  v <- newArray (0,9) 0
  sequence [writeArray v i (i+1) | i <- [0..9]]
  return v

-- A generator to zero-out a vector
clear_vec :: (SymVec repr, SSym repr, LamPure repr, SymLoop repr,
      Applicative m) =>
     J m (HV h repr) Int
     -> J m (HV h repr) (IOUArray Int Int)
     -> J m (HV h repr) (IO ())
clear_vec n a = 
 loop_ (int 0) (n +: int (-1)) (int 1) (lam $ \i ->
    vec_set $$ weakens a $$ var i $$ int 0)

-- The zero-out function to use now (at the present stage)
tclr :: IOUArray Int Int -> IO ()
tclr = runRI (lam (\v -> clear_vec (int 5) (var v)))

-- The code to use later
tclc = 
  "\\x_0 -> TSCore.loop_ 0 ((GHC.Num.+) 5 (-1)) 1 (\\x_1 -> "++
  "TSCore.vec_set_ x_0 x_1 0)" ==
   (runCI (lam (\v -> clear_vec (int 5) (var v))))

-- The imperative addition and multiplication
infixl 6 +::
(+::) :: (SSym repr, SymBind repr, Monad m) => 
	 repr (m Int) -> repr (m Int) -> repr (m Int)
x +:: y = gret add $$$ x $$$ y

infixl 7 *::
(*::) :: (SSym repr, SymBind repr, Monad m) => 
	 repr (m Int) -> repr (m Int) -> repr (m Int)
x *:: y = gret mul $$$ x $$$ y


-- Matrix-vector multiplication; the dimensions n and m are assumed
-- known statically (the same assumption was made om talk-problems.ml)
-- We use weakens all throughout (we need a signature though)
-- We assume the input vector is long.
mvmul0 :: (SymMat repr, SymVec repr, SymBind repr, SSym repr, LamPure repr,
	   SymLoop repr, Applicative m) =>
     Int -> Int
     -> J m (HV h repr) (IOUArray (Int, Int) Int)
     -> J m (HV h repr) (IOUArray Int Int)
     -> J m (HV h repr) (IOUArray Int Int)
     -> J m (HV h repr) (IO ())
mvmul0 n m a v v' =
 clear_vec (int n) v' >>:
 loop_ (int 0) (int (m-1)) (int 1) (lam $ \j ->
  loop_ (int 0) (int (n -1)) (int 1) (lam $ \i ->
   (vec_set $$ weakens v' $$ weakens (var i)) =<<:
    (vec_get $$ weakens v' $$ weakens (var i)) +::
     (mat_get $$ weakens a $$ weakens (var i) $$ weakens (var j)) *::
     (vec_get $$ weakens v $$ weakens (var j))
  ))

tms0c = "\\x_0 -> \\x_1 -> \\x_2 -> "++
  "(GHC.Base.>>) "++
  "(TSCore.loop_ 0 ((GHC.Num.+) 5 (-1)) 1 (\\x_3 -> "++
    "TSCore.vec_set_ x_2 x_3 0)) "++
  "(TSCore.loop_ 0 9 1 (\\x_3 -> "++
   "TSCore.loop_ 0 4 1 (\\x_4 -> "++
     "(GHC.Base.>>=) "++
     "(Control.Monad.ap (Control.Monad.ap (GHC.Base.return (GHC.Num.+)) "++
       "(TSCore.vec_get_ x_2 x_4)) "++
     "(Control.Monad.ap (Control.Monad.ap (GHC.Base.return (GHC.Num.*)) "++
       "(TSCore.mat_get_ x_0 x_4 x_3)) (TSCore.vec_get_ x_1 x_3))) "++
    "(TSCore.vec_set_ x_2 x_4))))"
 ==
 (runCI (lam $ \a -> (lam $ \vin -> (lam $ \v' ->
   mvmul0 5 10 (weakens (var a)) (weakens (var vin)) (var v')))))

-- Run the code now, to test it
tms0r = do
  a  <- sample_a
  v  <- sample_v
  v' <- newArray (0,4) 100
  runRI (lam $ \a -> (lam $ \vin -> (lam $ \v' ->
    mvmul0 5 10 (weakens (var a)) (weakens (var vin)) (var v'))))
   a v v'
  getAssocs v' >>= print  
-- [(0,330),(1,385),(2,440),(3,495),(4,550)]

-- Strip-mining, split-factor b, statically known
loop_nested :: (SSym repr, LamPure repr, SymLoop repr, SymMin repr,
		Applicative m) =>
     Int
     -> Int
     -> Int
     -> J m (HV h repr) (Int -> IO ())
     -> J m (HV h repr) (IO ())
loop_nested b lb ub body =
    loop_ (int lb) (int ub) (int b) (lam $ \ii ->
     loop_ (var ii) (min_ (var ii +: int (b-1)) (int ub)) (int 1)
      (weakens body))


-- The body remains the same; the use of weakens is crucial!

mvmul1 :: (SymMat repr, SymVec repr, SymBind repr, SSym repr, LamPure repr,
	   SymMin repr,
	   SymLoop repr, Applicative m) =>
     Int -> Int -> Int
     -> J m (HV h repr) (IOUArray (Int, Int) Int)
     -> J m (HV h repr) (IOUArray Int Int)
     -> J m (HV h repr) (IOUArray Int Int)
     -> J m (HV h repr) (IO ())
mvmul1 b n m a v v' =
 clear_vec (int n) v' >>:
 loop_nested b 0 (m-1) (lam $ \j ->
  loop_nested b 0 (n-1) (lam $ \i ->
   (vec_set $$ weakens v' $$ weakens (var i)) =<<:
    (vec_get $$ weakens v' $$ weakens (var i)) +::
     (mat_get $$ weakens a $$ weakens (var i) $$ weakens (var j)) *::
     (vec_get $$ weakens v $$ weakens (var j))
  ))


tms1c = "\\x_0 -> \\x_1 -> \\x_2 -> "++
 "(GHC.Base.>>) "++
 "(TSCore.loop_ 0 ((GHC.Num.+) 5 (-1)) 1 (\\x_3 -> "++
   "TSCore.vec_set_ x_2 x_3 0)) "++
 "(TSCore.loop_ 0 9 2 (\\x_3 -> "++
  "TSCore.loop_ x_3 (GHC.Classes.min ((GHC.Num.+) x_3 1) 9) 1 (\\x_4 -> "++
  "TSCore.loop_ 0 4 2 (\\x_5 -> "++
  "TSCore.loop_ x_5 (GHC.Classes.min ((GHC.Num.+) x_5 1) 4) 1 (\\x_6 -> "++
  "(GHC.Base.>>=) (Control.Monad.ap (Control.Monad.ap (GHC.Base.return (GHC.Num.+)) (TSCore.vec_get_ x_2 x_6)) (Control.Monad.ap (Control.Monad.ap (GHC.Base.return (GHC.Num.*)) (TSCore.mat_get_ x_0 x_6 x_4)) (TSCore.vec_get_ x_1 x_4))) (TSCore.vec_set_ x_2 x_6))))))"
 ==
 (runCI (lam $ \a -> (lam $ \v -> (lam $ \v' ->
   mvmul1 2 5 10 (weakens (var a)) (weakens (var v)) (var v')))))

-- run it for testing
tms1r = do
  a  <- sample_a
  v  <- sample_v
  v' <- newArray (0,4) 100
  runRI (lam $ \a -> (lam $ \v -> (lam $ \v' ->
    mvmul1 2 5 10 (weakens (var a)) (weakens (var v)) (var v'))))
   a v v'
  getAssocs v' >>= 
   return .  (== [(0,330),(1,385),(2,440),(3,495),(4,550)])

-- A loop-insertion primitive
insloop :: (SSym repr, SymLoop repr, LamPure repr, Applicative m) =>
     J (CPSA (HV hw repr (IO ())) m) (HV hw repr) Int ->
     J (CPSA (HV hw repr (IO ())) m) (HV hw repr) Int ->
     J (CPSA (HV hw repr (IO ())) m) (HV hw repr) Int ->
     J (CPSA (HV hw repr (IO ())) m) (HV ha repr) Int
insloop lb ub step = J (CPSA (\k -> 
  (unCPSA . unJ  $ (loop_ (weaken lb) (weaken ub) (weaken step) $ lam (\x -> 
   		 J (CPSA (\k0 -> 
		   throw k0 $
                   fmap (\kr -> J(\ (h1,(hw,h)) -> unJ (kr h1 hw) h))
		        (k (pure (\h1 _ -> J(\h -> unJ x (h1,undefined))))))))
		       ))
        (\x -> fmap (\xv h1 hw -> J(\h -> unJ (xv h1 hw) (hw,h))) x)
  ))

tk1 = "\\x_0 -> TSCore.loop_ 1 2 3 (\\x_1 -> let z_2 = x_1\n"++
       "                                     in TSCore.mat_set_ x_0 z_2 z_2 0)"
 ==
 (runCI (runJCPSA 
  (lam $ \a ->
    resetJ $ let_ (insloop (int 1) (int 2) (int 3))
             (\i -> mat_set $$ (weaken (var a)) $$ var i $$ var i $$ int 0))))



-- Like sloop nested, but lift up the first loop
-- The blocking factor is static (as it should be)
-- If we want it to be `dynamic' (that is, int code)
-- we have to make the type higher-rank since b is used in
-- different contexts (in the inner and outer loops)
loop_nested_exch :: (SSym repr, SymLet repr, LamPure repr,
		     SymLoop repr, SymMin repr, Applicative m) =>
     Int -> Int -> Int ->
     J (CPSA (HV hw repr (IO ())) m) (HV ha repr) (Int -> IO ()) ->
     J (CPSA (HV hw repr (IO ())) m) (HV ha repr) (IO ())
loop_nested_exch b lb ub body =
    let_ (insloop (int lb) (int ub) (int b)) (\ii ->
     loop_ (var ii) (min_ (var ii +: int (b-1)) (int ub)) (int 1)
      (weakens body))


-- The signature tells all the features we use
mvmul2 :: (SymMat repr, SymVec repr, SymBind repr, SSym repr, SymLet repr, 
	   LamPure repr, SymLoop repr, SymMin repr, Applicative m) =>
     Int -> Int -> Int
     -> J (CPSA (HV h repr (IO ())) m) (HV h repr) (IOUArray (Int, Int) Int)
     -> J (CPSA (HV h repr (IO ())) m) (HV h repr) (IOUArray Int Int)
     -> J (CPSA (HV h repr (IO ())) m) (HV h repr) (IOUArray Int Int)
     -> J (CPSA (HV h repr (IO ())) m) (HV h repr) (IO ())
mvmul2 b n m a v v' =
 clear_vec (int n) v' >>:
 (resetJ $
 loop_nested_exch b 0 (m-1) (lam $ \j ->
  loop_nested_exch b 0 (n-1) (lam $ \i ->
   (vec_set $$ weakens v' $$ weakens (var i)) =<<:
    (vec_get $$ weakens v' $$ weakens (var i)) +::
     (mat_get $$ weakens a $$ weakens (var i) $$ weakens (var j)) *::
     (vec_get $$ weakens v $$ weakens (var j))
  )))

tms2c = "\\x_0 -> \\x_1 -> \\x_2 -> "++
  "(GHC.Base.>>) "++
  "(TSCore.loop_ 0 ((GHC.Num.+) 5 (-1)) 1 (\\x_3 -> "++
  "TSCore.vec_set_ x_2 x_3 0)) "++
  "(TSCore.loop_ 0 9 2 (\\x_3 -> "++
   "TSCore.loop_ 0 4 2 (\\x_4 -> let z_5 = x_3\n                                                                                                                                                                            in "++
   "TSCore.loop_ z_5 (GHC.Classes.min ((GHC.Num.+) z_5 1) 9) 1 (\\x_6 -> "++
   "let z_7 = x_4\n                                                                                                                                                                                                                                                    in "++
   "TSCore.loop_ z_7 (GHC.Classes.min ((GHC.Num.+) z_7 1) 4) 1 (\\x_8 -> "++
    "(GHC.Base.>>=) (Control.Monad.ap (Control.Monad.ap (GHC.Base.return (GHC.Num.+)) (TSCore.vec_get_ x_2 x_8)) (Control.Monad.ap (Control.Monad.ap (GHC.Base.return (GHC.Num.*)) (TSCore.mat_get_ x_0 x_8 x_6)) (TSCore.vec_get_ x_1 x_6))) (TSCore.vec_set_ x_2 x_8))))))"
  ==
 (runCI (runJCPSA (lam $ \a -> (lam $ \v -> (lam $ \v' ->
  resetJ $ mvmul2 2 5 10
    (weakens (var a)) (weakens (var v)) (var v'))))))

tms2r = do
  a  <- sample_a
  v  <- sample_v
  v' <- newArray (0,4) 100
  runRI (runJCPSA (lam $ \a -> (lam $ \v -> (lam $ \v' ->
    resetJ $ mvmul2 2 5 10 (weakens (var a)) (weakens (var v)) (var v')))))
   a v v'
  getAssocs v' >>= 
   return .  (== [(0,330),(1,385),(2,440),(3,495),(4,550)])


main = sequence [return tclc, return tms0c, return tms1c, tms1r, 
		 return tk1,
		 return tms2c, tms2r] 
        >>= print . and