{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, TypeSynonymInstances, 
NoMonomorphismRestriction, FlexibleContexts, 
FunctionalDependencies, ScopedTypeVariables, 
TypeOperators #-}

{- Code / Term Interpretations for Data and Loop -}

module Code (Code(..), MCode, MCode', term) where

import StateCont
import Control.Monad
import Control.Monad.Identity
import Prelude hiding (div)
import Type
import Term
import Struct
import Data
import Control
import Array


{- Generate Term assembly code.

   To use, pass an expression that produces a single monadic value m
   (r t) to `termExpr', or a function to `termFun'.

-}  
 
-- Code is parameterized with a phantom type to implement Data.  We
-- represent that type info as data in Term, making use of typeOf.
-- This keeps the Term data type simple, no need for GADTs etc..  
data Code t = Code { unCode :: Term }


-- CompState passed around in the SC monad: bindings and var numbering.
data CompState = CompState Int

-- Monad used for Code interpretation.  This is a state continuation
-- monad with monadic type (Data return type) t, and compilation
-- result r.  Below, r is either TermExpr for expressions or Fun for
-- functions.
type MCode'  = SC CompState Term
type MCode t = MCode' t  

{- Wrapping it would make some things easier and some things harder..
data MCode t = MCode (SC Compstate Term)
instance Monad MCode where
  return v = MCode $ return v
  (>>=) (MCode sc) f = MCode (sc >> f') where
    f' v = sc where
      (MCode sc) = f v
-}

-- All generated variables are numbers prefixed with this.
regPrefix = "t"
regName n = regPrefix ++ show n

-- Generate some variables from the pool.  For Let insertion and
-- Lambda compilation.
nameList :: String -> Int -> MCode [String]
nameList base need = SC c where
  c (CompState n) k = k (CompState n') names where
    n' = n + need
    names = map ((base ++) . show) [n..(n+need-1)]

getNameCounter = SC c where
  c (CompState n) k = k (CompState n) n
  
setNameCounter n = SC c where
  c (CompState _) k = k (CompState n) ()


-- Insert a Let binding around the current context, and insert a
-- reference in the current expression hole.  This acts as a prompt:
-- control can't pass outside of the embedding Let, protecting the
-- context in which the variable reference is valid.

-- Pass explicit type tag.
mVarTyped typ term@(Code sterm) = SC c where  
  c (CompState n) k = 
    Let var sterm (k state' ref) where
      state' = CompState $ n+1
      ref = Code $ Ref $ var
      var = Var typ (regName n)

-- Reify Haskell type information as type tag.
mVar :: TypeOf (Code t) => (Code t) -> MCode (Code t)
mVar term = mVarTyped (typeOf term) term
      
-- Statement sequencing works just like Let, but ignores the term
-- return value so doesn't need to introduce a variable.
mVoid :: (Code t) -> MCode (Code ())
mVoid (Code sterm) = 
  SC (\s k -> Begin sterm (k s (Code Void)))

-- Structure unpacking is similar to Let.  Here mStruct is used to
-- generate a list of variables stored in the Unpack binding form, and
-- a structured Code object that propagates those variables as a Code
-- object with the proper structure.

-- TODO: mStruct could be updated to take termMatched as an argument
-- to avoid generating variables for terms that are already just
-- references.  Currently we leave this form of duplication to a
-- downstream optimizer, i.e. GCC or LLVM.  Maybe not worth
-- bothering..
mUnpack rs@(Code termPacked) = do
  (Code term, code) <- mStruct "s" $ Just $ (unrepStruct rs)
  let 
    Just vars   = varTree term
    termMatched = termTree vars termPacked -- match structure
    in do
    SC (\s k -> Unpack vars termMatched (k s code))
  

-- Structured variable generation.  This generates both a structure of
-- Code representations for injection into type code and an untyped
-- Term equivalent expressed as Atom / Cons / Nil of Ref.
instance DataWord t => StructVar (L (Code t)) where  
  structVar = sv where
    -- Reuse variable references when possible.
    -- sv _ id (Just t@(L (Code (Ref _)))) = (id, t)
    sv prefix id _ = (id + 1, L (Code $ term)) where
      typ = typeOf (undefined :: Code t)
      term = Ref $ Var typ $ prefix ++ show id


-- Generate Term Var and Code representation for a structure type.
-- This finds the size of the ras type, constructs a list of variable
-- names, plugs them into a Code representation and extracts the
-- variable references from that representation.  The term behind the
-- Maybe argument is for avoiding introduction of variable to variable
-- bindings.
mStruct :: forall s sr. 
           (Struct Code s sr,
            StructVar sr)
           => String -> Maybe sr -> MCode (Code s, sr)
-- sr : struct of representations
-- rs : representation of struct
mStruct namePrefix sr0 = do
  n <- getNameCounter
  let (n', sr :: sr) = structVar namePrefix n sr0
      rs = repStruct sr
    in do
    setNameCounter n'
    return $ (rs, sr)
  

-- Compile a monadic term forking the current compilation state.  This
-- delimits the context for Let insertion, which will manipulate the
-- toplevel continuation so needs to be localized.
--   sub:   monadic sub-term
--   code:  wrapper for the result     

-- Note: state will be forked!  The Let insertion mechanism doesn't
-- allow recovery of state.
mBlock :: MCode (Code t') -> MCode (Code t)
mBlock sub = SC main where
  main s k = k s $ Code $ (subTerm s sub)


cl t = mVar . Code . (Lit t) . show
cop2 typ opc (Code a) (Code b) = mVar (Code $ Op typ opc [a,b])
type TypedOp2 r t t' = String -> (Code t) -> (Code t) -> MCode (Code t')

icop2 = cop2 (Type AInt 0)   :: TypedOp2 r Tint Tint
fcop2 = cop2 (Type AFloat 0) :: TypedOp2 r Tfloat Tfloat

ibcop2 = cop2 (Type ABool 0) :: TypedOp2 r Tint Tbool
fbcop2 = cop2 (Type ABool 0) :: TypedOp2 r Tfloat Tbool

-------- Data (MCode r) Code
cconv typ nam (Code t) = mVar (Code $ Op typ nam [t])
instance Data MCode' Code where  
  true  = cl (Type ABool 0) True
  false = cl (Type ABool 0) False
  f2i = cconv (Type AInt 0)   "f2i"
  i2f = cconv (Type AFloat 0) "i2f"
  unpack = mUnpack
  pack ras = mVar $ packed where
    packed = Code $ Pack as
    (Code as) = repStruct ras

          
instance DataRing MCode' Code Tfloat where  
  add = fcop2 "add"
  sub = fcop2 "sub"
  mul = fcop2 "mul"
  eq  = fbcop2 "eq"
  lt  = fbcop2 "lt"
  lit = cl (Type AFloat 0)
  
instance DataField MCode' Code Tfloat where  
  div = fcop2 "div"
  
instance DataRing MCode' Code Tint where  
  add = icop2 "add"
  sub = icop2 "sub"
  mul = icop2 "mul"
  eq  = ibcop2 "eq"
  lt  = ibcop2 "lt"
  lit = cl (Type AInt 0)

-- The state that's passed around in the state-continuation (SC) monad
-- is an environment of variable bindings.
  


-- Evaluate the monadic term by passing initial environment and a
-- continuation.  Note that the return type of the continuation (Term)
-- is fixed by mVar.
subTerm s (SC c) = c s $ \s (Code t) -> t

-- Toplevel term, start with initial compilation state.
term = subTerm (CompState 0)


-- Bridge Code representation of structuring and (Haskell,
-- metalanguage) structuring of Code representations.

-- This collapses (Car|Cdr (Cons . .)) and Unatom (Atom .) for virtual
-- (= compile time) structs.

_atom   (L (Code t))    = Code (Atom t)
_unatom (Code (Atom t)) = (L (Code t))
_unatom (Code t)        = (L (Code (Unatom t)))

_nil   ()         = (Code Nil)
_unnil (Code Nil) = ()

_cons   (Code t1, Code t2)  = (Code (Cons t1 t2))
_uncons (Code (Cons t1 t2)) = (Code t1, Code t2)
_uncons (Code t)            = (Code (Car t), Code (Cdr t))
                              


-- Virtual data structures.  Used by SArray
instance StructRepr Code where
  
  rep0   = _nil
  rep1   = _atom
  rep2   = _cons
  
  unrep0 = _unnil
  unrep1 = _unatom
  unrep2 = _uncons

  -- Pack/unpack primitive data structures.  
          
unRef (Ref v) = v         

instance Repr Code

-- instance loopvar m r a => typeof (code a)
instance Control MCode' Code where

        
  ifte (Code c) mx my = do
    (Code x) <- mBlock mx
    (Code y) <- mBlock my
    return $ Code $ If c x y

  ret (Code t) = return $ Code $ Ret t

  letrec openDef openExpr = do
    [fName] <- nameList "fun" 1                  -- Name comes from the MCode monad's state.
    let ref         = Var (Type AVoid 0) fName   -- Create a symbolic reference Term.
        codeRef     = Code $ Ref ref             -- Generate a typed Code reference to pass to open forms.
        mDef        = openDef  codeRef           -- Close definition ..
        mExpr       = openExpr codeRef           -- .. and expression
        cLetRec d t = Code $ LetRec [(ref, d)] t -- wrap reference, definition and expr terms.
      in do
      (Code def) <- mDef  -- Lambda is already properly delimited, so just bind here.
      (Code expr) <- mBlock mExpr -- Evaluate mExpr in a delimited context for mVar.
      return $ Code $ LetRec [(ref, def)] expr
  
  def name mterm = do
    ct@(Code term) <- mterm
    mVoid $ Code $ Topdef (Var (typeOf ct) name) term
    
  -- Get [Var] and Code representations of arguments. Plug Code rep
  -- into HOS and perform code generation in an isolated lexical
  -- context block.  Connect args and body in final Lambda term.
  lam hos = do
    (Code argTerm, argCode) <- mStruct "a" Nothing
    Just argVars <- return $ varTree argTerm
    (Code body) <- mBlock $ hos argCode
    return $ Code $ Lambda argVars body
        
  -- Unpack the Code representation of the arguments into a list of
  -- variables and wrap in App.
  app (Code f) ras =
    return $ Code $ App f as where
      (Code as) = repStruct ras
    
    
-- The type tag of the result expression (array element) is obtained
-- from the array's type, which needs to be at least order 1.
derefType (Type _ 0) = error $ "Dereferencing order 0 Var."    
derefType (Type base order) = Type base (order - 1)
mArrVar a = mVarTyped (derefType $ aType a)


-- Arrays are either plain variables or stored in a struct and
-- accessed through Car / Cdr.  In order to find the element type we
-- need to recurse down a struct tree if there is one.
aType (Ref (Var typ _)) = typ
aType (Unatom t) = typ where Type (ATree (AAtom typ))   0 = aType t
aType (Car t)    = typ where Type (ATree (ACons typ _)) 0 = aType t
aType (Cdr t)    = typ where Type (ATree (ACons _ typ)) 0 = aType t
aType t = error $ "aType: " ++ show t

instance DataWord t => Array MCode' Code ((->) Tint) t where
  get (Code a) (Code i) = mArrVar a (Code $ Get a i)
  set (Code a) (Code i) (Code e) = mVoid $ Code $ Set a i e