{-# LANGUAGE RankNTypes, TemplateHaskell #-}
{-# LANGUAGE FlexibleContexts, FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}

-- Tagless Final with staging
-- The core of the code generation framework
-- This file hides several implementation types and functions,
-- exporting only the safe operations 

module TSCore (
  SSym(..),
  LamPure(..),
  J(..),
  runCS, unR,

  HV,
  H,			-- data constructor not exported!
  hmap,
  weaken,  (+:), (*:), lam, var,
  Extends(..),

  runR, runRI,
  runC, runCI,
  AssertPos, assertPos,

  SymDIV, (/:),
  SymLet, let_,

  -- Generating imperative code
  SymBind, SymLoop, gret, (>>:), ($$$), (=<<:), SymMat(..), SymVec(..),
  loop_,
  SymMin, min_,

  -- Needed only for well-labellness tests, not used in production
  LamLPure, State, liftState, Label, lamL, runCL, runCLI

  ) where

import Control.Applicative
import Control.Exception (assert)
import Language.Haskell.TH
import Language.Haskell.TH.Syntax
import Language.Haskell.TH.Ppr
import Data.List (isSuffixOf)

import Control.Monad (ap)
import Data.Array.MArray
import Data.Array.IO


-- The object language (EDSL)

-- A Haskell value of the type SSym repr => repr a 
-- represents an expression in the object language of the type a

-- There is no lam form; See the class LamPure below
-- dedicated to abstractions.
class SSym repr where
    int :: Int -> repr Int
    add :: repr (Int -> Int -> Int)
    mul :: repr (Int -> Int -> Int)

    ($$) :: repr (a->b) -> (repr a -> repr b)

infixl 2 $$

{-
The connection to MetaOCaml

Unlike MetaOCaml, based on brackets and escapes for building
code values, we use combinators for code generation (cf. Xi and Thiemann).
Here is how our combinators could be defined in terms of
brackets and escapes:

let int x  =  <x>
let add    =  <(*)>
let x $$ y =  <~x ~y>

let lam f   =  <fun x -> ~(f <x>)>

The combinators int, add, ($$) build code _values_, which are inert
at present stage and whose building involves no effect.
-}

exS1 :: SSym repr => repr Int
exS1 = (add $$ int 1) $$ int 2
exS1c = runCS exS1 

-- ------------------------------------------------------------------------
-- The interpreters for our DSL

-- The evaluator
newtype R a = R{unR :: a}

instance SSym R where
    int = R
    add = R (+)
    mul = R (*)
    R x $$ R y = R $ x y

-- The C interpreter, the compiler
-- We mostly use it to show the code

type VarCounter = Int		-- we will see the need for it shortly, lamS

			-- the pure code value, the analogue
			-- of <1> or <fun x -> x> in MetaOCaml
newtype C a = C{unC :: VarCounter -> Exp}

instance SSym C where
    int = C . const . LitE . IntegerL . fromIntegral
    add = C $ const $ VarE $ '(Prelude.+)
    mul = C $ const $ VarE $ '(Prelude.*)
    C x $$ C y = C $ \vc -> AppE (x vc) (y vc)

runCS :: C a -> String
runCS m = pprint $ unC m 0

-- ------------------------------------------------------------------------
-- Adding `purely generated' lambda-expressions to our EDSL
-- Any effect of generating their body can't propagate past
-- the binder.
-- See below for more general lambda-expressions


class LamPure repr where
    lamS :: (repr a -> repr b) -> repr (a->b)

-- We extend the R and C interpreters
-- The extensions are quite like
-- that in our Tagless-Final paper

instance LamPure R where
    lamS f = R $ unR . f . R

instance LamPure C where
    lamS f = C $ \vc -> 
	     let name = mkName $ "x_" ++ show vc
	         body = unC (f (C . const $ VarE name)) (succ vc)
             in LamE [VarP name] body

-- ------------------------------------------------------------------------
-- Applicative transformers
-- They are much simpler than monad transformers since applicatives
-- are composable. We introduce them as a syntactic sugar to be
-- used shortly

-- The newtype wrapper (J m repr a) denotes the composition
-- of the applicative m with a member of SSym. The latter
-- isn't quite applicative, but close (it doesn't
-- have the polymorphic `pure')

newtype J m repr a = J{unJ :: m (repr a)}

-- The composition of an applicative with a member of SSym
-- is itself a member of SSym

instance (Applicative m, SSym repr) => SSym (J m repr) where
    int = J . pure . int
    add = J (pure add)
    mul = J (pure mul)
    J x $$ J y = J (($$) <$> x <*> y)

-- ------------------------------------------------------------------------
-- Hypothetical code values
-- A Haskell value of the type SSym repr => HV h repr a 
-- represents an expression in the object language of the type a
-- in the (generator) environment h
-- Recall that (->) h is an applicative.

type HV h = J ((->) h)

-- HV values are thus themselves members of SSym, as we have
-- just seen

hmap :: (h2 -> h1) -> HV h1 repr a -> HV h2 repr a
hmap f e = J $ \h2 -> unJ e (f h2)

runH :: HV () repr a -> repr a
runH m = unJ m ()


{-
Like MetaOCaml, we also keep track of the (future stage) environment,
distinguishing a closed expression like <1> from the open expression
<x>.  We keep a far detailed track than does MetaOCaml, tracking each
free variable rather than a set of variables of the same classifier.
Thus our type (HV h repr a) is _roughly_ equivalent to (('h,'a) code) of
MetaOCaml.  Roughly because 'h' has a lot more structure than just the
single qualifier.
-}

-- The data constructor H and the elimination for href are
-- NOT exported!

-- A component of the environment h
newtype H r s a = H (r a) 			-- just to attach s

href :: HV (H repr s a,h) repr a		-- hypothesis reference
href = J $ \ (H x,h) -> x



-- We can define the following, but it doesn't help us much...
-- We have do define lam anew
lamH :: (SSym repr, LamPure repr) =>
       (forall s. HV (H repr s a,h) repr a -> HV (H repr s a,h) repr b)
       ->  HV h repr (a->b)
lamH f = J $ \h -> lamS (\x -> (unJ (f href)) (H x,h))


-- Effectful code generator
-- A Haskell value of the type
-- (Applicative m, SSym repr) => J m (HV h repr) a
-- represents a generator, in the generating applicative m,
-- of the expression in the object language of the type a
-- in the (generator) environment h

{-
We are more precise than MetaOCaml in another aspect, keeping track
of effects of code generation as well. In MetaOCaml, as in OCaml in
general, all expressions are effectful. Therefore, our type 
(J m (HV h repr) a) corresponds to a general MetaOCaml expression, with
effects, that will generate the code of type a in the future
stage environment h.
-}

weaken :: Applicative m => J m (HV h repr) a -> J m (HV (h',h) repr) a
weaken (J m) = J (fmap (hmap snd) m)

-- Convenience functions

infixl 7 *:
infixl 6 +:

(+:) :: (SSym repr) => repr Int -> repr Int -> repr Int
x +: y = add $$ x $$ y

(*:) :: (SSym repr) => repr Int -> repr Int -> repr Int
x *: y = mul $$ x $$ y


-- We use a higher-order abstract syntax. The body of lam receives
-- an HV value rather than (m (HV ...)) value: that is, the future-stage
-- variable is a pure code value, cannot have any effects.
-- That becomes important when we get to delimited continuations,
-- since pure values are answer-type--polymorphic.

lam :: (Functor m, SSym repr, LamPure repr) =>
       (forall s. HV (H repr s a,h) repr a -> J m (HV (H repr s a,h) repr) b)
       -> J m (HV h repr) (a->b)
lam f = J (fmap (\body -> J $ \h -> lamS (\x -> unJ body (H x,h))) 
                (unJ (f href)))

var :: (Applicative m, SSym repr) => HV h repr a -> J m (HV h repr) a
var = J . pure


-- Generic weakening
-- The constraint Extends h h'  asserts that the environment h'
-- is either equal to h or is an extension of h with more bindings.
-- The assertion is proved by inserting the needed number of
-- weaken

class Extends h h' where
    weakens :: Applicative m => J m (HV h repr) a -> J m (HV h' repr) a

instance Extends h h where
    weakens = id

-- The following is simple but obviously non-generic.
-- See the Regions paper for the generic, inductive weakening
instance Extends h (h1,h) where
    weakens = weaken

instance Extends h (h2,(h1,h)) where
    weakens = weaken . weaken

instance Extends h (h3,(h2,(h1,h))) where
    weakens = weaken . weaken . weaken


-- we could compose weakens with var, so var' will have a signature
--   var' :: Extends h h' => HV h repr a -> m (HV h' repr a)
-- That will be more convenient. Let be explicit with weaken for now.


-- Functions to obtain the results

runR :: Applicative i => J i (HV () R) a -> i a
runR = fmap (unR . runH) . unJ

-- Specialized instances of runR
runRI :: J Identity (HV () R) a -> a
runRI = runIdentity . runR

runC :: Applicative i => J i (HV () C) a -> i String
runC = fmap (runCS . runH) . unJ

-- Specialized instances of runC
runCI :: J Identity (HV () C) a -> String
runCI = runIdentity . runC


-- ------------------------------------------------------------------------
-- Adding a new DSL feature: assertions
-- (to check the assertion insertion)
--   assertPos test x 
-- checks that the value of test is positive. If so, it returns x.
-- Otherwise, a run-time error is generated

class AssertPos repr where
    assertPos :: repr Int -> repr a -> repr a


instance AssertPos R where
    assertPos (R test) (R x) = R $ assert (test > 0) x

instance AssertPos C where
    assertPos (C test) (C x) = 
	C $ \vc -> AppE (AppE assertName (testE (test vc))) (x vc)
     where
     testE exp = InfixE (Just exp) ge (Just zero)
     zero = LitE . IntegerL . fromIntegral $ 0
     ge = VarE $ '(Prelude.>)
     assertName = VarE $ 'Control.Exception.assert


instance (Applicative m, AssertPos repr) => AssertPos (J m repr) where
    assertPos (J x) (J y) = J (assertPos <$> x <*> y)


-- ------------------------------------------------------------------------
-- Adding a new DSL feature: integer division
-- (for the sake of the assertion-insertion example)

class SymDIV repr where
    divS :: repr (Int -> Int -> Int)

instance SymDIV R where
    divS = R div

instance SymDIV C where
    divS = C $ const $ VarE $ 'Prelude.div

instance (Applicative m, SymDIV repr) => SymDIV (J m repr) where
    divS = J (pure divS)

infixl 7 /:

(/:) :: (Applicative m, SSym repr, SymDIV repr) => 
	J m (HV h repr) Int -> J m (HV h repr) Int -> J m (HV h repr) Int
x /: y = divS $$ x $$ y


-- ------------------------------------------------------------------------
-- Adding a new DSL feature: let-expressions
-- Not surprisingly, it looks like a `composition' of lam and app

class SymLet repr where
    let_S :: repr a -> (repr a -> repr b) -> repr b

-- In a pure language, let_S is a reverse application, as expected
instance SymLet R where
    let_S rhs f = f rhs

instance SymLet C where
    let_S rhs f = C $ \vc -> 
	     let name = mkName $ "z_" ++ show vc
	         body = unC (f (C . const $ VarE name)) (succ vc)
             in LetE [ValD (VarP name) (NormalB $ unC rhs vc) []] body

-- We define let_ as a `combination' of lam and application.
-- Hence the following type, which ensures that the let-bound variable
-- does not escape the scope of the binding.
let_ :: (SSym repr, SymLet repr, Applicative m) => 
      J m (HV h repr) a
      -> (forall s. HV (H repr s a,h1) repr a -> J m (HV (H repr s a,h) repr) b)
      -> J m (HV h repr) b
let_ rhs fbody = J $ comb <$> (unJ rhs) <*> (unJ (fbody href))
 where comb rhsv bv = J $ \h -> let_S (unJ rhsv h) (\x -> unJ bv (H x,h))


-- ------------------------------------------------------------------------
-- Adding a new DSL feature: labels on lam and variables
-- and checking of the labels

type Label = Int

class LamLPure repr where
    lamSL :: Label -> (repr a -> repr b) -> repr (a->b)
    check_label :: Label -> repr a -> repr a    

instance LamLPure C where
    lamSL l f = C $ \vc -> 
	     let name = mkName $ "x_" ++ show vc ++ "_" ++ show l
	         body = unC (f (C . const $ VarE name)) (succ vc)
             in LamE [VarP name] body

    check_label l v@(C c) = 
         let (VarE name) = c 0 in
         if isSuffixOf ("_"++show l) (show name) then v
            else error $ unwords ["Ill-labeled: got varname",show name,
				  " but expected the label", show l]

-- used internally
hrefL :: LamLPure repr => Label -> HV (H repr s a,h) repr a
hrefL l = J $ \ (H x,h) -> check_label l x

-- lamL uniquely labels the introduced free variable promise and the
-- generated binder
lamL :: (Functor m, Monad m, SSym repr, LamLPure repr) =>
       (forall s. HV (H repr s a,h) repr a 
	-> J (State Label m) (HV (H repr s a,h) repr) b)
       -> J (State Label m) (HV h repr) (a->b)
lamL f = J (do
	    l <- newLabel
	    fmap (\body -> J $ \h -> lamSL l (\x -> unJ body (H x,h)))
	                (unJ (f (hrefL l))))

runCL :: (Functor m, Monad m) => J (State Label m) (HV () C) a -> m String
runCL m = fmap (runCS . runH) $ (evalState (unJ m) 0)

-- Specialized instances of runCL
runCLI :: J (State Label Identity) (HV () C) a -> String
runCLI = runIdentity . runCL


-- ------------------------------------------------------------------------
-- We extend our DSL to generate imperative code

-- We add higher-order constants for bind and return
class SymBind repr where
    ret_   :: Monad m => repr (a -> m a)
    bind_  :: Monad m => repr (m a -> (a -> m b) -> m b)
    app_   :: Monad m => repr (m (a->b) -> m a -> m b)
    scolon :: Monad m => repr (m () -> m a -> m a)

instance SymBind R where
    ret_   = R return
    bind_  = R (>>=)
    app_   = R (ap)
    scolon = R (>>)

instance SymBind C where
    ret_   = C (const (VarE 'return))
    bind_  = C (const (VarE '(>>=)))
    app_   = C (const (VarE 'ap))
    scolon = C (const (VarE '(>>)))

instance (Applicative m, SymBind repr) => SymBind (J m repr) where
    ret_   = J (pure ret_)
    bind_  = J (pure bind_)
    app_   = J (pure app_)
    scolon = J (pure scolon)
           

-- A bit of syntactic sugar
gret :: (SSym repr, SymBind repr, Monad m) => repr a -> repr (m a)
gret a = ret_ $$ a

-- sequencing of two computations: `the semi-colon'
infixl 1 >>:
(>>:) :: (SSym repr, SymBind repr, Monad m) =>
	 repr (m ()) -> repr (m a) -> repr (m a)
x >>: y = scolon $$ x $$ y

-- pure-function application
infixl 2 $$$
($$$) :: (SSym repr, SymBind repr, Monad m) => 
	 repr (m (a -> b)) -> repr (m a) -> repr (m b)
x $$$ y = app_ $$ x $$ y

-- call-by-value (monadic) application
infixr 1 =<<:
(=<<:) :: (SSym repr, SymBind repr, Monad m) =>
	  repr (a -> m b) -> repr (m a) -> repr (m b)
x =<<: y = bind_ $$ y $$ x

-- Matrices with 0-based indices
mat_new_ :: MArray IOUArray e IO =>
	   Int -> Int -> e -> IO (IOUArray (Int,Int) e)
mat_new_ i j e = newArray ((0,i),(0,j)) e

mat_get_ :: MArray IOUArray e IO =>
	   IOUArray (Int,Int) e -> Int -> Int -> IO e
mat_get_ a i j = readArray a (i,j)

mat_set_ :: MArray IOUArray e IO =>
	   IOUArray (Int,Int) e -> Int -> Int -> e -> IO ()
mat_set_ a i j e = writeArray a (i,j) e


class SymMat repr where
    mat_new :: MArray IOUArray e IO =>
	       repr (Int -> Int -> e -> IO (IOUArray (Int,Int) e))
    mat_get :: MArray IOUArray e IO =>
	       repr (IOUArray (Int,Int) e -> Int -> Int -> IO e)

    mat_set :: MArray IOUArray e IO =>
	       repr (IOUArray (Int,Int) e -> Int -> Int -> e -> IO ())


instance SymMat R where
    mat_new = R mat_new_
    mat_get = R mat_get_
    mat_set = R mat_set_

instance SymMat C where
    mat_new = C (const (VarE 'mat_new_))
    mat_get = C (const (VarE 'mat_get_))
    mat_set = C (const (VarE 'mat_set_))

instance (Applicative m, SymMat repr) => SymMat (J m repr) where
    mat_new = J (pure mat_new)
    mat_get = J (pure mat_get)
    mat_set = J (pure mat_set)
           

-- Vectors with 0-based indices
vec_new_ :: MArray IOUArray e IO =>
	   Int -> e -> IO (IOUArray Int e)
vec_new_ i e = newArray (0,i) e

vec_get_ :: MArray IOUArray e IO =>
	   IOUArray Int e -> Int -> IO e
vec_get_ a i = readArray a i

vec_set_ :: MArray IOUArray e IO =>
	   IOUArray Int e -> Int -> e -> IO ()
vec_set_ a i e = writeArray a i e

class SymVec repr where
    vec_new :: MArray IOUArray e IO =>
	       repr (Int -> e -> IO (IOUArray Int e))
    vec_get :: MArray IOUArray e IO =>
	       repr (IOUArray Int e -> Int -> IO e)
    vec_set :: MArray IOUArray e IO =>
	       repr (IOUArray Int e -> Int -> e -> IO ())


instance SymVec R where
    vec_new = R vec_new_
    vec_get = R vec_get_
    vec_set = R vec_set_

instance SymVec C where
    vec_new = C (const (VarE 'vec_new_))
    vec_get = C (const (VarE 'vec_get_))
    vec_set = C (const (VarE 'vec_set_))

instance (Applicative m, SymVec repr) => SymVec (J m repr) where
    vec_new = J (pure vec_new)
    vec_get = J (pure vec_get)
    vec_set = J (pure vec_set)


-- The loop combinator, a higher-order constant
--     sloop lb up step (\x -> body)

-- imperative loop with a step
loop_comb :: Int -> Int -> Int -> (Int -> IO ()) -> IO ()
loop_comb lb ub _ _ | lb > ub = return ()
loop_comb lb ub step body = body lb >> loop_comb (lb + step) ub step body
  

class SymLoop repr where
    loop_combS :: repr (Int -> Int -> Int -> (Int -> IO ()) -> IO ())


instance SymLoop R where
    loop_combS = R loop_comb

instance SymLoop C where
    loop_combS = C (const (VarE 'loop_))

instance (Applicative m, SymLoop repr) => SymLoop (J m repr) where
    loop_combS = J (pure loop_combS)
           

-- Syntactic sugar
loop_ :: (SSym repr, SymLoop repr) =>
    repr Int -> repr Int -> repr Int -> repr (Int -> IO ()) ->
         repr (IO ())
loop_ lb ub step body = loop_combS $$ lb $$ ub $$ step $$ body


-- The min function
class SymMin repr where
    minS :: repr (Int -> Int -> Int)

instance SymMin R where
    minS = R min

instance SymMin C where
    minS  = C (const . VarE $ 'min)

instance (Applicative m, SymMin repr) => SymMin (J m repr) where
    minS = J (pure minS)


-- Syntactic sugar
min_ :: (SSym repr, SymMin repr) =>
    repr Int -> repr Int -> repr Int
min_ x y = minS $$ x $$ y


-- ------------------------------------------------------------------------
-- Various pieces missing from the standard libraries 
-- We define them here

-- There is no Applicative Identity in standard libraries

newtype Identity a = Identity{runIdentity :: a}
instance Functor Identity where
    fmap f = Identity . f . runIdentity
instance Applicative Identity where
    pure = Identity
    f <*> x = Identity $ (runIdentity f) (runIdentity x)

instance Monad Identity where
    return  = Identity
    x >>= f = f (runIdentity x)

-- Before 7.04, the standard library misses the following
-- Uncomment the following for older versions of GHC
{-
instance Applicative (Either String) where
    pure = Right
    Left e <*> _  = Left e
    _ <*> Left e  = Left e
    Right f <*> Right m = Right (f m)
-}

-- The state Applicative and monad
newtype State l m a = State{unState :: l -> m (a,l)}

instance Monad m => Functor (State l m) where
    fmap f m = pure f <*> m

instance Monad m => Applicative (State l m) where
    pure    = return
    f <*> m = do
	      fv <- f
	      mv <- m
	      return (fv mv)

instance Monad m => Monad (State l m) where
    return x = State $ \l -> return (x,l)
    m >>= f  = State $ \l -> do (mv,l1) <- unState m l
                                unState (f mv) l1
    
newLabel :: (Monad m, Num l) => State l m l
newLabel = State $ \l -> return (l, l+1)

evalState :: Monad m => State l m a -> l -> m a
evalState (State m) l = m l >>= return . fst

liftState :: Monad m => m a -> State l m a
liftState a = State $ \l -> a >>= \v -> return (v,l)