DesugarSyntax.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE StandaloneDeriving,TypeSynonymInstances,FlexibleInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE LambdaCase #-}

-- Transform all type variables to unique rigid type variables.
module DesugarSyntax where

import Control.Monad
import Control.Monad.State
import Control.Monad.Identity
import qualified Data.Map.Strict as M

import Syntax
import FreshNames

type Desugar = StateT DState (FreshMT Identity)

type IdTVMap = M.Map Id (VType Desugared)

type IdAbModMap = M.Map Id (AbMod Desugared)

data DState = MkDState { env :: IdTVMap
                       , abModEnv :: IdAbModMap }

getDState :: Desugar DState
getDState = get

putDState :: DState -> Desugar ()
putDState = put

freshRTVar :: Id -> Desugared -> Desugar (VType Desugared)
freshRTVar x a = do n <- fresh
                    return $ RTVar (x ++ "$r" ++ (show n)) a

freshRigid :: Id -> Desugar Id
freshRigid x = do n <- fresh
                  return $ x ++ "$r" ++ show n

pullId :: VType Desugared -> Id
pullId (RTVar id _) = id
pullId _ = error "pullId called on something other than a rigid tvar"

putEnv :: IdTVMap -> Desugar ()
putEnv p = do s <- getDState
              putDState $ s { env = p }

getEnv :: Desugar IdTVMap
getEnv = do s <- getDState
            return $ env s

putAbModEnv :: IdAbModMap -> Desugar ()
putAbModEnv p = do s <- getDState
                   putDState $ s { abModEnv = p }

getAbModEnv :: Desugar IdAbModMap
getAbModEnv = do s <- getDState
                 return $ abModEnv s


initDState :: DState
initDState = MkDState M.empty M.empty

desugar :: Prog Refined -> Prog Desugared
desugar (MkProg xs) = MkProg $ evalFresh m
  where m = evalStateT (mapM desugarTopTm' xs) initDState
        desugarTopTm' tm =
          do putEnv $ M.empty -- purge mappings from previous context.
             desugarTopTm tm

-- no explicit refinements
desugarTopTm :: TopTm Refined -> Desugar (TopTm Desugared)
desugarTopTm (DataTm dt a) = DataTm <$> desugarDataT dt <*> pure (refToDesug a)
desugarTopTm (ItfTm itf a) = ItfTm <$> desugarItf itf <*> pure (refToDesug a)
desugarTopTm (DefTm def a) = DefTm <$> desugarMHDef def <*> pure (refToDesug a)

-- explicit refinements:
-- + type variables get fresh ids
desugarDataT :: DataT Refined -> Desugar (DataT Desugared)
desugarDataT (DT dt ps                ctrs a) = do
  -- id val & eff ty vars constructors
  -- generate fresh ids for ty vars
  xs' <- mapM (freshRigid . fst) ps
  -- map old value type variables to fresh ones
  putEnv $ M.fromList [(x, RTVar x' a') | ((x, VT), x') <- zip ps xs']
  -- map old effect type variables to fresh ones
  putAbModEnv $ M.fromList [(x, AbRVar x' a') | ((x, ET), x') <- zip ps xs']
  -- [(new var name, val or eff)]
  let ps' = [(x', k) | ((_, k), x') <- zip ps xs']
   -- the following will only use but not modify the DState
  ctrs' <- mapM desugarCtr ctrs
  -- ps' contains new fresh names
  return $ DT dt ps' ctrs' a'
  where a' = refToDesug a

-- explicit refinements:
-- + type variables get fresh ids
desugarItf :: Itf Refined -> Desugar (Itf Desugared)
desugarItf (Itf itf ps cmds a) = do
  -- generate fresh ids for type variables
  xs' <- mapM (freshRigid . fst) ps
  let env' = M.fromList [(x, RTVar x' a') | ((x, VT), x') <- zip ps xs']       -- map old value type variables to fresh ones
  let abModEnv' = M.fromList [(x, AbRVar x' a') | ((x, ET), x') <- zip ps xs'] -- map old effect type variables to fresh ones
  -- [(new var name, val or eff)]
  let ps' = [(x', k) | ((_, k), x') <- zip ps xs']

  -- before desugaring each cmd, we need to reset the env
  cmds' <- mapM (\c -> do putEnv env'
                          putAbModEnv abModEnv'
                          desugarCmd c)
                cmds
  -- reset afterwards, too
  putEnv env'
  putAbModEnv abModEnv'
  return $ Itf itf ps' cmds' a'
  where a' = refToDesug a

-- no explicit refinements
desugarMHDef :: MHDef Refined -> Desugar (MHDef Desugared)
desugarMHDef (Def hdr ty xs a) = do
  ty' <- desugarCType ty
  xs' <- mapM desugarClause xs
  return $ Def hdr ty' xs' (refToDesug a)

-- explicit refinements:
-- + type variables get fresh ids
desugarCmd :: Cmd Refined -> Desugar (Cmd Desugared)
desugarCmd (Cmd cmd  ps                 xs       y a) = do
  --  id   val & eff ty vars  arg tys  return ty
  -- generate fresh ids for type variables
  ps' <- mapM (freshRigid . fst) ps
  -- map old value type variables to fresh ones
  putEnv $ M.fromList [(p, RTVar p' a') | ((p, VT), p') <- zip ps ps']
  -- map old effect type variables to fresh ones
  putAbModEnv $ M.fromList [(p, AbRVar p' a') | ((p, ET), p') <- zip ps ps']
  -- [(new var name, val or eff)]
  let ps'' = [(p', k) | ((_, k), p') <- zip ps ps']
  xs' <- mapM desugarVType xs
  y' <- desugarVType y
  return $ Cmd cmd ps'' xs' y' a'
  where a' = refToDesug a

-- no explicit refinements
desugarCtr :: Ctr Refined -> Desugar (Ctr Desugared)
desugarCtr (Ctr ctr xs a) = do
  xs' <- mapM desugarVType xs
  return $ Ctr ctr xs' (refToDesug a)

-- explicit refinements:
-- + replace val ty vars by corresponding MkRTVar's of env,
--   generate new fresh one and add if not in env yet
-- + replace 'String' by 'List Char'
desugarVType :: VType Refined -> Desugar (VType Desugared)
desugarVType (TVar x a) =
  do env <- getEnv
     case M.lookup x env of
       Nothing -> do ty <- freshRTVar x (refToDesug a)
                     putEnv $ M.insert x ty env
                     return ty
       Just ty -> return ty
desugarVType (DTTy dt ts a) = do ts' <- mapM desugarTyArg ts
                                 return $ DTTy dt ts' (refToDesug a)
desugarVType (SCTy ty a) = do ty' <- desugarCType ty
                              return $ SCTy ty' (refToDesug a)
-- replace 'String' by 'List Char'
desugarVType (StringTy a) = return $ desugaredStrTy (refToDesug a)
desugarVType (IntTy a) = return $ IntTy (refToDesug a)
desugarVType (CharTy a) = return $ CharTy (refToDesug a)

-- nothing happens on this level
desugarTyArg :: TyArg Refined -> Desugar (TyArg Desugared)
desugarTyArg (VArg t a) = VArg <$> desugarVType t <*> pure (refToDesug a)
desugarTyArg (EArg ab a) = EArg <$> desugarAb ab <*> pure (refToDesug a)

-- nothing happens on this level
desugarCType :: CType Refined -> Desugar (CType Desugared)
desugarCType (CType ports peg a) =
  CType <$> mapM desugarPort ports <*> desugarPeg peg <*> pure (refToDesug a)

-- nothing happens on this level
desugarPort :: Port Refined -> Desugar (Port Desugared)
desugarPort (Port adj ty a) = Port <$> desugarAdj adj <*> desugarVType ty <*> pure (refToDesug a)

-- nothing happens on this level
desugarPeg :: Peg Refined -> Desugar (Peg Desugared)
desugarPeg (Peg ab ty a) = Peg <$> desugarAb ab <*> desugarVType ty <*> pure (refToDesug a)

-- nothing happens on this level
desugarAb :: Ab Refined -> Desugar (Ab Desugared)
desugarAb (Ab v (ItfMap m b) a) = do v' <- desugarAbMod v
                                     m' <- mapM (mapM desugarTyArg) m
                                     return $ Ab v' (ItfMap m' (refToDesug b)) (refToDesug a)

-- explicit desugaring:
-- + replace effect type variables by corresponding MkAbRVar's of abModEnv,
--   generate new fresh one and add if not in env yet
desugarAbMod :: AbMod Refined -> Desugar (AbMod Desugared)
desugarAbMod (EmpAb a) = return $ EmpAb (refToDesug a)
desugarAbMod (AbVar x a) =
  do env <- getAbModEnv
     case M.lookup x env of
       Nothing -> do n <- fresh
                     let var = AbRVar (x ++ "$r" ++ (show n)) (refToDesug a)
                     putAbModEnv $ M.insert x var env
                     return var
       Just var -> return var

-- nothing happens on this level
desugarAdj :: Adj Refined -> Desugar (Adj Desugared)
desugarAdj (Adj (ItfMap m b) a) = do m' <- mapM (mapM desugarTyArg) m
                                     return $ Adj (ItfMap m' (refToDesug b)) (refToDesug a)

-- no explicit desugaring: clauses (and constituents) unaffected between Refine/Desugar phase
desugarClause :: Clause Refined -> Desugar (Clause Desugared)
desugarClause (Cls ps tm a) = do ps' <- mapM desugarPattern ps
                                 tm' <- desugarTm tm
                                 return $ Cls ps' tm' (refToDesug a)

desugarTm :: Tm Refined -> Desugar (Tm Desugared)
desugarTm (SC x a) = SC <$> desugarSComp x <*> pure (refToDesug a)
desugarTm (StrTm s a) = return $ StrTm s (refToDesug a)
desugarTm (IntTm n a) = return $ IntTm n (refToDesug a)
desugarTm (CharTm c a) = return $ CharTm c (refToDesug a)
desugarTm (TmSeq tm1 tm2 a) = TmSeq <$> desugarTm tm1 <*> desugarTm tm2 <*> pure (refToDesug a)
desugarTm (Use u a) = Use <$> desugarUse u <*> pure (refToDesug a)
desugarTm (DCon d a) = DCon <$> desugarDCon d <*> pure (refToDesug a)

desugarPattern :: Pattern Refined -> Desugar (Pattern Desugared)
desugarPattern (VPat v a) = VPat <$> desugarVPat v <*> pure (refToDesug a)
desugarPattern (CmdPat c vs k a) = do vs' <- mapM desugarVPat vs
                                      return $ CmdPat c vs' k (refToDesug a)
desugarPattern (ThkPat x a) = return $ ThkPat x (refToDesug a)

desugarVPat :: ValuePat Refined -> Desugar (ValuePat Desugared)
desugarVPat (VarPat x a) = return $ VarPat x (refToDesug a)
desugarVPat (DataPat x xs a) = DataPat x <$> mapM desugarVPat xs <*> pure (refToDesug a)
desugarVPat (IntPat i a) = return $ IntPat i (refToDesug a)
desugarVPat (CharPat c a) = return $ CharPat c (refToDesug a)
desugarVPat (StrPat s a) = return $ StrPat s (refToDesug a)


desugarSComp :: SComp Refined -> Desugar (SComp Desugared)
desugarSComp (SComp xs a) = SComp <$> mapM desugarClause xs <*> pure (refToDesug a)

desugarUse :: Use Refined -> Desugar (Use Desugared)
desugarUse (App use xs a) = App <$> desugarUse use <*> mapM desugarTm xs <*> pure (refToDesug a)
desugarUse (Op op a) = Op <$> desugarOperator op <*> pure (refToDesug a)

desugarOperator :: Operator Refined -> Desugar (Operator Desugared)
desugarOperator (Mono x a) = return $ Mono x (refToDesug a)
desugarOperator (Poly x a) = return $ Poly x (refToDesug a)
desugarOperator (CmdId x a) = return $ CmdId x (refToDesug a)

desugarDCon :: DataCon Refined -> Desugar (DataCon Desugared)
desugarDCon (DataCon id xs a) = DataCon id <$> mapM desugarTm xs <*> pure (refToDesug a)

-- Helpers

refToDesug :: Refined -> Desugared
refToDesug (Refined loc) = Desugared loc