content: {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE DeriveAnyClass #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# OPTIONS_GHC -fno-warn-name-shadowing #-}
  {-# OPTIONS_GHC -fno-warn-orphans #-}
  {-# OPTIONS_GHC -fdefer-typed-holes #-}
  {-# OPTIONS_GHC -fno-warn-unused-do-bind #-}
  
  -- |
  
  module Gibbon.L2.Typecheck
      ( tcExp, tcProg, TCError(..)
      , RegionSet(..)
      , LocationTypeState(..)
      , ConstraintSet(..)
      , LocConstraint(..)
      , Aliased, TcM )
      where
  
  import           Control.DeepSeq
  import           Control.Monad
  import           Control.Monad.Except
  import           Data.Foldable ( foldlM, foldrM )
  import qualified Data.Set as S
  import qualified Data.List as L
  import qualified Data.Map as M
  import qualified Safe as Sf
  
  import           Data.Maybe
  import           Text.PrettyPrint.GenericPretty
  
  import           Gibbon.Common
  import           Gibbon.L2.Syntax as L2
  -- import qualified Gibbon.L1.Syntax as L1
  
  -- | Constraints on locations.  Used during typechecking.  Roughly analogous to LocExp.
  data LocConstraint = StartOfC LocVar Region -- ^ Location is equal to start of this region.
                     -- Can't attach haddocks to data constructor arguments with < GHC 8.4.2
                     -- See https://github.com/haskell/haddock/pull/709.
                     | AfterConstantC Int     -- Number of bytes after.
                                      LocVar  -- Location which is before
                                      LocVar  -- Location which is after
                     | AfterVariableC Var     -- Name of variable v. This loc is size(v) bytes after.
                                      LocVar  -- Location which is before
                                      LocVar  -- Location which is after
                     | InRegionC LocVar Region -- Location is somewher within this region.
                     -- New constraints here for the SoA backend.
                     | GenSoALocC LocVar [((DataCon, FieldIndex), LocVar)]
                     | GetDataConLocSoAC LocVar 
                     | AssignC LocVar LocVar -- First location is equal to the second location.
                                             -- First location is rhs, second location is lhs
                     | GetFieldLocSoAC (DataCon, FieldIndex) LocVar
                     | FromEndC LocVar
  
    deriving (Read, Show, Eq, Ord, Generic, NFData, Out)
  
  
  -- | A set of constraints (which are re-used location expressions)
  -- which encode relationships between locations. These are used by the
  -- type checker to verify that locations are used correctly.
  --
  -- The following are valid constraints:
  --  * StartOfC LocVar Region
  --  * AfterConstantC Int LocVar LocVar
  --  * AfterVariableC Var LocVar LocVar
  --  * InRegionC LocVar Region
  --  * FromEndC
  --
  -- While the first four can appear in syntax before RouteEnds, the fifth
  -- (fromEnd) should only be introduced by the RouteEnds pass.
  newtype ConstraintSet = ConstraintSet { constraintSet :: S.Set LocConstraint }
    deriving (Read, Show, Eq, Ord, Generic, NFData, Out)
  
  -- | A location has been aliased if we have taken an offset of it while introducing a new
  -- location. These show up in the LocationTypeState below.
  type Aliased = Bool
  
  -- | A map from locations to their type state. This tracks the current state of a location
  -- at a certain point in a program. The type checker threads this through by taking an
  -- input map that represents the location states before checking a particular expression,
  -- and returns an output map that represents the location states after checking that
  -- particular expression.
  -- All we need to know in this map is whether each location is an input or output location,
  -- and whether we've taken an offset of the location in a letloc form.
  newtype LocationTypeState = LocationTypeState
      {
        tsmap :: M.Map LocVar (Modality,Aliased)
        -- ^ Each location has a modality and may have had an offset taken.
      }
      deriving (Read,Show,Eq,Ord, Generic, NFData, Out)
  
  -- | A region set is (as you would expect) a set of regions. They are the
  -- regions that are currently live while checking a particular expression.
  newtype RegionSet = RegionSet { regSet :: S.Set Region }
    deriving (Read, Show, Eq, Ord, Generic, NFData)
  
  
  -- | Shorthand for located expressions
  type Exp = Exp2
  
  -- | These are the kinds of errors that the type checker may throw. There are a
  -- few named errors that I thought would be common, like variables not being
  -- found, locations being used incorrectly, etc. For other errors, there's a
  -- GenericTC form that takes some expression and string
  -- (containing an error message).
  data TCError = GenericTC String Exp
               | VarNotFoundTC Var Exp
               | UnsupportedExpTC Exp
               | LocationTC String Exp LocVar LocVar
               | ModalityTC String Exp LocVar LocationTypeState
                 deriving (Read,Eq,Ord,Generic,NFData)
  
  instance Show TCError where
      show (GenericTC str e) = "Error typechecking L2 Program\nIn the expression:\n" ++ (sdoc e) ++ "\n" ++ str ++ "\n"
      show (VarNotFoundTC v e) = "Variable not found: " ++ (show v) ++ "\nIn the expression:\n" ++ (sdoc e) ++ "\n"
      show (UnsupportedExpTC e) = "Unsupported expression:\n" ++ (sdoc e) ++ "\n"
      show (LocationTC str e lv1 lv2) = "Location typechecking error: " ++ str ++ "\nIn the expression:\n" ++ (sdoc e)
                                        ++ "\nLocations: " ++ (show lv1) ++ ", " ++ (show lv2) ++ "\n"
      show (ModalityTC str e lv lts) = "Modality typechecking error: " ++ str ++ "\nIn the expression:\n" ++ (sdoc e)
                                       ++ "\nLocation: " ++ (show lv) ++ "\nLocation type state: " ++ (show lts) ++ "\n"
  
  -- | The type checking monad. Just for throwing errors, but could in the future be parameterized
  --   by whatever other logging, etc, monad we want the compiler to use.
  type TcM a = (Except TCError) a
  
  
  
  -- | Check an expression. Given the data definitions, an general type environment, a function map,
  --   a constraint set, a region set, an (input) location state map, and the expression, this function
  --   will either throw an error, or return a pair of expression type and new location state map.
  tcExp :: DDefs Ty2 -> Env2 FreeVarsTy Ty2 -> FunDefs2
        -> ConstraintSet -> RegionSet -> LocationTypeState -> Exp
        -> TcM (Ty2, LocationTypeState)
  tcExp ddfs env funs constrs regs tstatein exp =
  
      case exp of
  
        VarE v ->
            -- Look up a variable in the environment
            do ty <- lookupVar env (fromVarToFreeVarsTy v) exp
               return (ty, tstatein)
  
        LitE _i -> return (IntTy, tstatein)
  
        CharE _i -> return (CharTy, tstatein)
  
        FloatE _i -> return (FloatTy, tstatein)
  
        LitSymE _v -> return (SymTy, tstatein)
  
        AppE v ls args ->
            -- Checking function application involves a few steps:
            --  (1) We need to make sure the inputs/ouptuts line up with the expected
            --      types for the function.
            --  (2) We need to update the location state map with information about what output
            --      locations have been written to by the called function and must now be input
            --      locations.
            --  (3) We need to make sure that if we pass a packed structure as an argument, its
            --      location is among the passed-in locations.
            do let (ArrowTy2 locVars arrIns _arrEffs arrOut _locRets _isPar) =
                       case M.lookup v funs of
                         Just f -> funTy f
                         Nothing -> error $ "tcExp: Unbound function: " ++ sdoc v
  
               -- Get types of arguments
               (in_tys, tstate) <- foldlM
                                     (\(tys, st) a -> do
                                           (ty, st') <- recur st a
                                           pure (tys ++ [ty], st'))
                                     ([],tstatein) args
  
               -- Check arity
               if (length args) /= (length in_tys)
                 then throwError $ GenericTC ("Arity mismatch. Expected:" ++ show (length in_tys) ++ " Got:" ++ show (length args)) exp
                 else pure ()
  
               -- (1)
               mapM (uncurry $ ensureEqualTyNoLoc exp) (fragileZip in_tys arrIns)
  
               -- (3) Check location of argument
               let tyls = concatMap locsInTy in_tys
               case L.find (\loc -> not $ S.member loc (S.fromList ls)) tyls of
                 Nothing -> return ()
                 Just not_in_ls -> throwError $ GenericTC ("Packed argument location expected: " ++ show not_in_ls) exp
               let handleTS ts (l,Output) =  switchOutLoc exp ts l
                   handleTS ts _ = return ts
               -- (2)
               tstate' <- foldM handleTS tstate $ zip ls $ L.map (\(LRM _ _ m) -> m) locVars
               -- Use locVars used at call-site in the returned type
               let arrOutMp = M.fromList $ zip (L.map (\(LRM l _ _) -> l) locVars) ls
                   arrOut'  = substLoc arrOutMp arrOut
  
               return (arrOut',tstate')
  
        PrimAppE pr es -> do
                 -- Special case because we can't lookup the type of the function pointer
                 let es' = case pr of
                             VSortP{} -> init es
                             InplaceVSortP{} -> init es
                             _        -> es
                 (tys,tstate) <- tcExps ddfs env funs constrs regs tstatein es'
  
                 -- Pattern matches would be one way to check length safely, but then the
                 -- error would not go through our monad:
                 let len2 = checkLen exp pr 2 es
                     len1 = checkLen exp pr 1 es
                     len0 = checkLen exp pr 0 es
                     len3 = checkLen exp pr 3 es
                     len4 = checkLen exp pr 4 es
  
                     mk_bools = do
                       len0
                       pure (BoolTy, tstate)
  
                     bool_ops = do
                       len2
                       _ <- ensureEqualTy (es !! 0) BoolTy (tys !! 0)
                       _ <- ensureEqualTy (es !! 1) BoolTy (tys !! 1)
                       pure (BoolTy, tstate)
  
                     int_ops = do
                       len2
                       _ <- ensureEqualTy (es !! 0) IntTy (tys !! 0)
                       _ <- ensureEqualTy (es !! 1) IntTy (tys !! 1)
                       pure (IntTy, tstate)
  
                     float_ops = do
                       len2
                       _ <- ensureEqualTy (es !! 0) FloatTy (tys !! 0)
                       _ <- ensureEqualTy (es !! 1) FloatTy (tys !! 1)
                       pure (FloatTy, tstate)
  
                     int_cmps = do
                       len2
                       _ <- ensureEqualTy (es !! 0) IntTy (tys !! 0)
                       _ <- ensureEqualTy (es !! 1) IntTy (tys !! 1)
                       pure (BoolTy, tstate)
  
                     float_cmps = do
                       len2
                       _ <- ensureEqualTy (es !! 0) FloatTy (tys !! 0)
                       _ <- ensureEqualTy (es !! 1) FloatTy (tys !! 1)
                       pure (BoolTy, tstate)
  
                     char_cmps = do
                       len2
                       _ <- ensureEqualTy (es !! 0) CharTy (tys !! 0)
                       _ <- ensureEqualTy (es !! 1) CharTy (tys !! 1)
                       pure (BoolTy, tstate)
  
                 case pr of
                   MkTrue  -> mk_bools
                   MkFalse -> mk_bools
                   AddP    -> int_ops
                   SubP    -> int_ops
                   MulP    -> int_ops
                   DivP    -> int_ops
                   ModP    -> int_ops
                   ExpP    -> int_ops
                   FAddP   -> float_ops
                   FSubP   -> float_ops
                   FMulP   -> float_ops
                   FDivP   -> float_ops
                   FExpP   -> float_ops
                   EqIntP  -> int_cmps
                   LtP     -> int_cmps
                   GtP     -> int_cmps
                   LtEqP   -> int_cmps
                   GtEqP   -> int_cmps
                   EqFloatP -> float_cmps
                   EqCharP  -> char_cmps
                   FLtP     -> float_cmps
                   FGtP     -> float_cmps
                   FLtEqP   -> float_cmps
                   FGtEqP   -> float_cmps
                   OrP     -> bool_ops
                   AndP    -> bool_ops
  
                   RandP -> return (IntTy, tstate)
                   FRandP -> return (FloatTy, tstate)
  
                   FloatToIntP -> do
                     len1
                     ensureEqualTy exp FloatTy (tys !! 0)
                     return (IntTy, tstate)
  
                   IntToFloatP -> do
                     len1
                     ensureEqualTy exp IntTy (tys !! 0)
                     return (FloatTy, tstate)
  
                   FSqrtP -> do
                     len1
                     ensureEqualTy exp FloatTy (tys !! 0)
                     return (FloatTy, tstate)
  
                   FTanP -> do
                     len1
                     ensureEqualTy exp FloatTy (tys !! 0)
                     return (FloatTy, tstate)
  
                   Gensym -> len0 >>= \_ -> pure (SymTy, tstate)
  
                   EqSymP -> do
                     len2
                     ensureEqualTy exp SymTy (tys !! 0)
                     ensureEqualTy exp SymTy (tys !! 1)
                     return (BoolTy,tstate)
  
                   EqBenchProgP _ -> do
                     len0
                     return (BoolTy,tstate)
  
                   DictEmptyP ty -> do
                     len1
                     let [a] = tys
                     _ <- ensureEqualTy exp ArenaTy a
                     case es !! 0 of
                       (VarE var) ->
                           do ensureArenaScope exp env (Just var)
                              return (SymDictTy (Just var) (stripTyLocs ty), tstate)
                       _ -> throwError $ GenericTC "Expected arena variable argument" exp
  
                   DictInsertP ty -> do
                     len4
                     let [a,d,k,v]  = tys
                     _ <- ensureEqualTy exp ArenaTy a
                     _ <- ensureEqualTy exp SymTy k
                     _ <- ensureEqualTyNoLoc exp ty v
                     case d of
                       SymDictTy ar _ty ->
                           case es !! 0 of
                             (VarE var) ->
                                 do ensureArenaScope exp env ar
                                    ensureArenaScope exp env (Just var)
                                    return (SymDictTy (Just var) (stripTyLocs ty), tstate)
                             _ -> throwError $ GenericTC "Expected arena variable argument" exp
                       _ -> throwError $ GenericTC "Expected SymDictTy" exp
  
                   DictLookupP ty -> do
                     len2
                     let [d,k]  = tys
                     case d of
                       SymDictTy ar _ty ->
                           do _ <- ensureEqualTy exp SymTy k
                              ensureArenaScope exp env ar
                              return (ty, tstate)
                       _ -> throwError $ GenericTC "Expected SymDictTy" exp
  
                   DictHasKeyP _ty -> do
                     len2
                     let [d,k]  = tys
                     case d of
                       SymDictTy ar _ty -> do _ <- ensureEqualTy exp SymTy k
                                              ensureArenaScope exp env ar
                                              return (BoolTy, tstate)
                       _ -> throwError $ GenericTC "Expected SymDictTy" exp
  
                   SizeParam -> do
                     len0
                     return (IntTy, tstate)
  
                   IsBig -> do
                     len2
                     let [ity, ety] = tys
                     ensureEqualTy exp ity IntTy
                     if isPackedTy ety
                     then pure (BoolTy, tstate)
                     else error "L1.Typecheck: IsBig expects a Packed value."
  
                   ErrorP _str ty -> do
                     len0
                     return (ty, tstate)
  
                   ReadPackedFile _fp _tycon _reg ty -> do
                     len0
                     return (ty, tstate)
  
                   WritePackedFile _ ty
                     | PackedTy{} <- ty  -> do
                       len1
                       let [packed_ty] = tys
                       _ <- ensureEqualTy exp packed_ty ty
                       pure (ProdTy [], tstate)
                     | otherwise -> error $ "writePackedFile expects a packed type. Given" ++ sdoc ty
  
                   ReadArrayFile _ ty -> do
                     len0
                     if isValidListElemTy ty
                     then return (VectorTy ty, tstate)
                     else throwError $ GenericTC "Not a valid list type" exp
  
                   RequestSizeOf -> do
                     len1
                     case (es !! 0) of
                       VarE{} -> if isPackedTy (tys !! 0)
                                 then return (IntTy, tstate)
                                 else case (tys !! 0) of
                                        SymTy -> return (IntTy, tstate)
                                        IntTy -> return (IntTy, tstate)
                                        _ -> throwError $ GenericTC "Expected PackedTy" exp
                       _ -> throwError $ GenericTC "Expected a variable argument" exp
  
                   VAllocP elty -> do
                     len1
                     checkListElemTy elty
                     let [i] = tys
                     _ <- ensureEqualTy (es !! 0) IntTy i
                     pure (VectorTy elty, tstate)
  
                   VFreeP elty -> do
                     len1
                     checkListElemTy elty
                     let [i] = tys
                     _ <- ensureEqualTy (es !! 0) (VectorTy elty) i
                     pure (ProdTy [], tstate)
  
                   VFree2P elty -> do
                     len1
                     checkListElemTy elty
                     let [i] = tys
                     _ <- ensureEqualTy (es !! 0) (VectorTy elty) i
                     pure (ProdTy [], tstate)
  
                   VLengthP elty -> do
                     let [ls] = tys
                     _ <- ensureEqualTy exp (VectorTy elty) ls
                     pure (IntTy, tstate)
  
                   VNthP elty -> do
                     let [ls, i] = tys
                     _ <- ensureEqualTy exp (VectorTy elty) ls
                     _ <- ensureEqualTy exp IntTy i
                     pure (elty, tstate)
  
                   VSliceP elty   -> do
                     let [from,to,ls] = tys
                     _ <- ensureEqualTy exp IntTy from
                     _ <- ensureEqualTy exp IntTy to
                     _ <- ensureEqualTy exp (VectorTy elty) ls
                     pure (VectorTy elty, tstate)
  
                   InplaceVUpdateP elty -> do
                     let [ls,i,val] = tys
                     _ <- ensureEqualTy exp (VectorTy elty) ls
                     _ <- ensureEqualTy exp IntTy i
                     _ <- ensureEqualTy exp elty val
                     pure (VectorTy elty, tstate)
  
                   VConcatP elty -> do
                     len1
                     let [ls] = tys
                     _ <- ensureEqualTy (es !! 0) (VectorTy (VectorTy elty)) ls
                     pure (VectorTy elty, tstate)
  
  
                   -- Given that the first argument is a list of type (VectorTy t),
                   -- ensure that the 2nd argument is function reference of type:
                   -- ty -> ty -> Bool
                   VSortP elty ->
                     case (es !! 1) of
                       VarE f -> do
                         len2
                         let [ls]   = tys
                             fn_ty  = lookupFEnvLocVar (fromVarToFreeVarsTy f) env
                             in_tys = inTys fn_ty
                             ret_ty = outTy fn_ty
                             err x  = throwError $ GenericTC ("vsort: Expected a sort function of type (ty -> ty -> Bool). Got"++ sdoc x) exp
                         _ <- ensureEqualTy (es !! 0) (VectorTy elty) ls
                         case in_tys of
                           [a,b] -> do
                              _ <- ensureEqualTy (es !! 1) a elty
                              _ <- ensureEqualTy (es !! 1) b elty
                              _ <- ensureEqualTy (es !! 1) ret_ty IntTy
                              pure (VectorTy elty, tstate)
                           _ -> err fn_ty
                       oth -> throwError $ GenericTC ("vsort: function pointer has to be a variable reference. Got"++ sdoc oth) exp
  
                   InplaceVSortP elty -> recur tstatein (PrimAppE (VSortP elty) es)
  
                   VMergeP elty -> do
                     len2
                     checkListElemTy elty
                     let [ls1,ls2] = tys
                     _ <- ensureEqualTy (es !! 0) (VectorTy elty) ls1
                     _ <- ensureEqualTy (es !! 1) (VectorTy elty) ls2
                     pure (VectorTy elty, tstate)
  
                   PDictInsertP kty vty -> do
                     len3
                     checkListElemTy kty
                     checkListElemTy vty
                     let [key, val, dict] = tys
                     _ <- ensureEqualTy (es !! 0) key kty
                     _ <- ensureEqualTy (es !! 1) val vty
                     _ <- ensureEqualTy (es !! 2) dict (PDictTy kty vty)
                     pure (PDictTy kty vty, tstate)
  
                   PDictLookupP kty vty -> do
                     len2
                     checkListElemTy kty
                     checkListElemTy vty
                     let [key, dict] = tys
                     _ <- ensureEqualTy (es !! 0) key kty
                     _ <- ensureEqualTy (es !! 1) dict (PDictTy kty vty)
                     pure (vty, tstate)
  
                   PDictAllocP kty vty -> do
                     len0
                     checkListElemTy kty
                     checkListElemTy vty
                     pure (PDictTy kty vty, tstate)
  
                   PDictHasKeyP kty vty -> do
                     len2
                     checkListElemTy kty
                     checkListElemTy vty
                     let [key, dict] = tys
                     _ <- ensureEqualTy (es !! 0) key kty
                     _ <- ensureEqualTy (es !! 1) dict (PDictTy kty vty)
                     pure (BoolTy, tstate)
  
                   PDictForkP kty vty -> do
                     len1
                     checkListElemTy kty
                     checkListElemTy vty
                     let [dict] = tys
                     _ <- ensureEqualTy (es !! 0) dict (PDictTy kty vty)
                     pure (ProdTy [PDictTy kty vty, PDictTy kty vty], tstate)
  
                   PDictJoinP kty vty -> do
                     len2
                     checkListElemTy kty
                     checkListElemTy vty
                     let [dict1, dict2] = tys
                     _ <- ensureEqualTy (es !! 0) dict1 (PDictTy kty vty)
                     _ <- ensureEqualTy (es !! 1) dict2 (PDictTy kty vty)
                     pure (PDictTy kty vty, tstate)
  
                   LLAllocP elty -> do
                     len0
                     checkListElemTy elty
                     pure (ListTy elty, tstate)
  
                   LLIsEmptyP elty -> do
                     len1
                     checkListElemTy elty
                     let [ll] = tys
                     _ <- ensureEqualTy (es !! 0) ll (ListTy elty)
                     pure (BoolTy, tstate)
  
                   LLConsP elty -> do
                     len2
                     checkListElemTy elty
                     let [elt, ll] = tys
                     _ <- ensureEqualTy (es !! 0) elt elty
                     _ <- ensureEqualTy (es !! 1) ll (ListTy elty)
                     pure (ListTy elty, tstate)
  
                   LLHeadP elty -> do
                     len1
                     checkListElemTy elty
                     let [ll] = tys
                     _ <- ensureEqualTy (es !! 0) ll (ListTy elty)
                     pure (elty, tstate)
  
                   LLTailP elty -> do
                     len1
                     checkListElemTy elty
                     let [ll] = tys
                     _ <- ensureEqualTy (es !! 0) ll (ListTy elty)
                     pure (ListTy elty, tstate)
  
                   LLFreeP elty -> do
                     len1
                     checkListElemTy elty
                     let [i] = tys
                     _ <- ensureEqualTy (es !! 0) (ListTy elty) i
                     pure (ProdTy [], tstate)
  
                   LLFree2P elty -> do
                     len1
                     checkListElemTy elty
                     let [i] = tys
                     _ <- ensureEqualTy (es !! 0) (ListTy elty) i
                     pure (ProdTy [], tstate)
  
                   LLCopyP elty -> do
                     len1
                     checkListElemTy elty
                     let [i] = tys
                     _ <- ensureEqualTy (es !! 0) (ListTy elty) i
                     pure (ListTy elty, tstate)
  
                   GetNumProcessors -> do
                     len0
                     pure (IntTy, tstate)
  
                   PrintInt -> do
                     len1
                     _ <- ensureEqualTy (es !!! 0) IntTy (tys !!! 0)
                     pure (ProdTy [], tstate)
  
                   PrintChar -> do
                     len1
                     _ <- ensureEqualTy (es !!! 0) CharTy (tys !!! 0)
                     pure (ProdTy [], tstate)
  
                   PrintFloat -> do
                     len1
                     _ <- ensureEqualTy (es !!! 0) FloatTy (tys !!! 0)
                     pure (ProdTy [], tstate)
  
                   PrintBool -> do
                     len1
                     _ <- ensureEqualTy (es !!! 0) BoolTy (tys !!! 0)
                     pure (ProdTy [], tstate)
  
                   PrintSym -> do
                     len1
                     _ <- ensureEqualTy (es !!! 0) SymTy (tys !!! 0)
                     pure (ProdTy [], tstate)
  
                   ReadInt  -> throwError $ GenericTC "ReadInt not handled" exp
  
                   SymSetEmpty -> do
                     len0
                     pure (SymSetTy, tstate)
  
                   SymSetInsert -> do
                     len2
                     _ <- ensureEqualTy (es !!! 0) SymSetTy (tys !!! 0)
                     _ <- ensureEqualTy (es !!! 1) SymTy (tys !!! 1)
                     pure (SymSetTy, tstate)
  
                   SymSetContains -> do
                     len2
                     _ <- ensureEqualTy (es !!! 0) SymSetTy (tys !!! 0)
                     _ <- ensureEqualTy (es !!! 1) SymTy (tys !!! 1)
                     pure (BoolTy, tstate)
  
                   SymHashEmpty -> do
                     len0
                     pure (SymHashTy, tstate)
  
                   SymHashInsert -> do
                     len3
                     _ <- ensureEqualTy (es !!! 0) SymHashTy (tys !!! 0)
                     _ <- ensureEqualTy (es !!! 1) SymTy (tys !!! 1)
                     _ <- ensureEqualTy (es !!! 2) SymTy (tys !!! 2)
                     pure (SymHashTy, tstate)
  
                   SymHashLookup -> do
                     len2
                     _ <- ensureEqualTy (es !!! 0) SymHashTy (tys !!! 0)
                     _ <- ensureEqualTy (es !!! 1) SymTy (tys !!! 1)
                     pure (SymTy, tstate)
  
                   SymHashContains -> do
                     len2
                     _ <- ensureEqualTy (es !!! 0) SymHashTy (tys !!! 0)
                     _ <- ensureEqualTy (es !!! 1) SymTy (tys !!! 1)
                     pure (BoolTy, tstate)
  
                   IntHashEmpty -> do
                     len0
                     pure (IntHashTy, tstate)
  
                   IntHashInsert -> do
                     len3
                     _ <- ensureEqualTy (es !!! 0) IntHashTy (tys !!! 0)
                     _ <- ensureEqualTy (es !!! 1) SymTy (tys !!! 1)
                     _ <- ensureEqualTy (es !!! 2) IntTy (tys !!! 2)
                     pure (IntHashTy, tstate)
  
                   IntHashLookup -> do
                     len2
                     _ <- ensureEqualTy (es !!! 0) IntHashTy (tys !!! 0)
                     _ <- ensureEqualTy (es !!! 1) SymTy (tys !!! 1)
                     pure (IntTy, tstate)
  
                   Write3dPpmFile{} -> throwError $ GenericTC "Write3dPpmFile not handled yet" exp
  
                   RequestEndOf{} -> throwError $ GenericTC  "tcExp of PrimAppE: RequestEndOf not handled yet" exp
  
        LetE (v, _ls, ty, e1@(AppE _f _ls1 _)) e2 -> do
          (ty1,tstate1) <- recur tstatein e1
          ensureEqualTyNoLoc exp ty1 ty
          let zipped = zip _ls _ls1
          (tstate', constrs') <- foldrM (\(l, l1) (st, c) -> do 
                                               r <- getRegion exp c l1
                                               let c' = extendConstrs (InRegionC l r) c
                                               pure (st, c') 
                           ) (tstate1, constrs) zipped  
          --r <- getRegion exp constrs _l1
          --let tstate1 = extendTS loc (Output,True) $ setAfter _l1 tstatein
          --let constrs1 = extendConstrs (InRegionC loc r) $ extendConstrs (FromEndC _l1) constrs
          
          let env' = extendVEnvLocVar (fromVarToFreeVarsTy v) ty env 
          tcExp ddfs env' funs constrs' regs tstate' e2
  
        LetE (v,_ls,ty,e1) e2 -> do
                 -- We get the type and new location state from e1
                 (ty1,tstate1) <- recur tstatein e1
                 ensureEqualTyNoLoc exp ty1 ty
                 let env' = extendVEnvLocVar (fromVarToFreeVarsTy v) ty env
                 -- Then we check e1 with that location state
                 tcExp ddfs env' funs constrs regs tstate1 e2
  
                   
                  
  
        IfE e1 e2 e3 -> do
  
                 -- Handle condition
                 (ty1,tstate1) <- recur tstatein e1
                 ensureEqualTyModCursor exp ty1 BoolTy
  
                 -- Check both branches
                 (ty2,tstate2) <- recur tstate1 e2
                 (ty3,tstate3) <- recur tstate1 e3
  
                 -- Combine the type states somehow (TODO: audit this)
                 tstate <- combineTStates exp tstate2 tstate3
  
                 ensureEqualTyModCursor exp ty2 ty3
                 return (ty2,tstate)
  
        MkProdE es -> do
                 (tys,tstate) <- tcExps ddfs env funs constrs regs tstatein es
                 return (ProdTy tys,tstate)
  
        ProjE i e -> do
  
                 (ty,tstate) <- recur tstatein e
                 tyi <- tcProj exp i ty
                 return (tyi, tstate)
  
        CaseE e brs -> do
  
                 (ty,tstate) <- recur tstatein e
                 case ty of
                   PackedTy _dc lin -> do
                           -- We need to know the "region" of all the vars (and locations) pattern matched by this case expresssion.
                           -- The "region" is the same as that of the case scrutinee.
                           reg <- getRegion e constrs lin
                           ensureMatchCases ddfs exp ty brs
                           (tys,tstate') <- tcCases ddfs env funs constrs regs tstate lin reg brs
                           foldM_ (ensureEqualTyModCursor exp) (tys !! 0) (Sf.tailErr tys)
                           return (tys !! 0,tstate')
                   _ -> error ("Expected packed type, got " ++ show ty)
  
        DataConE l dc es -> do
                 let dcty = getTyOfDataCon ddfs dc
                 (tys,tstate1) <- tcExps ddfs env funs constrs regs tstatein es
                 let args = lookupDataCon ddfs dc
  
                 if length args /= length es
                 then throwError $ GenericTC "Invalid argument length" exp
                 else do
                   sequence_ [ ensureEqualTyNoLoc exp ty1 ty2
                             | (ty1,ty2) <- zip args tys ]
                   -- -- TODO: need to fix this check for SoA location.
                   ensureDataCon exp dcty dc l tys constrs
                   tstate2 <- switchOutLoc exp tstate1 l
                   return (PackedTy dcty l, tstate2)
  
        TimeIt e _ty _b -> do
  
                 (ty1,tstate1) <- recur tstatein e
                 -- ensureEqualTy exp ty ty1
                 return (ty1,tstate1)
  
        SpawnE f locs args ->
          tcExp ddfs env funs constrs regs tstatein (AppE f locs args)
  
        SyncE -> pure (ProdTy [], tstatein)
  
        WithArenaE v e -> do
                let env' = extendVEnvLocVar (fromVarToFreeVarsTy v) ArenaTy env
                tcExp ddfs env' funs constrs regs tstatein e
  
        MapE _ _ -> throwError $ UnsupportedExpTC exp
  
        FoldE _ _ _ -> throwError $ UnsupportedExpTC exp
  
        Ext (LetRegionE r _ _ e) -> do
                 regs' <- regionInsert exp r regs
                 (ty,tstate) <- tcExp ddfs env funs constrs regs' tstatein e
                 return (ty,tstate)
  
        Ext (LetParRegionE r _ _ e) -> do
                 regs' <- regionInsert exp r regs
                 (ty,tstate) <- tcExp ddfs env funs constrs regs' tstatein e
                 return (ty,tstate)
        -- REMOVE THIS IDEALLY, we should not use this and there should be a polymorphic let Loc 
        --Ext (LetSoALocE loc e) -> do
        --  -- TODO: check why a ty of CursorTy is added, this type should change??
        --  let env' = extendVEnvLocVar loc CursorTy env
        --  let tstate1 = extendTS loc (Output, True) tstatein
        --  (ty,tstate2) <- tcExp ddfs env' funs constrs regs tstate1 e
        --  tstate3 <- removeLoc exp tstate2 loc
        --  return (ty,tstate3)
        Ext (LetLocE loc c e) -> do
                {- VS: TODO: bug, why is the type of loc always CursorTy? -}
                let env' = extendVEnvLocVar (fromLocVarToFreeVarsTy loc) CursorTy env
                case c of
                  GenSoALoc dcloc fieldLocs -> do
                    dcloc_reg <- getRegion exp constrs dcloc
                    let constrs' = extendConstrs (InRegionC (dcloc) dcloc_reg) constrs
                    --let nconstrs' = extendConstrs (InRegionC l r) $ nconstrs
                    (constrs'', fregs) <- foldrM (\(k, floc) (accC, accR) -> do
                                                                             freg <- getRegion exp accC floc
                                                                             return $ (extendConstrs (InRegionC (floc) freg) accC , (k, freg):accR)
                                           ) (constrs', []) fieldLocs
                    let new_reg = SoAR dcloc_reg fregs
                    let constrs''' = extendConstrs (InRegionC loc new_reg) $ constrs''
                    let env' = extendVEnvLocVar (fromLocVarToFreeVarsTy loc) CursorTy env
                    let tstate1 = extendTS loc (Output, True) tstatein
                    (ty,tstate2) <- tcExp ddfs env' funs constrs''' regs tstate1 e
                    tstate3 <- removeLoc exp tstate2 loc
                    return (ty,tstate3)
                  StartOfRegionLE r ->
                      do ensureRegion exp r regs
                         absentStart exp constrs r
                         let tstate1 = extendTS loc (Output,False) tstatein
                         let constrs1 = extendConstrs (StartOfC loc r) $ extendConstrs (InRegionC loc r) constrs
                         (ty,tstate2) <- tcExp ddfs env' funs constrs1 regs tstate1 e
                         tstate3 <- removeLoc exp tstate2 loc
                         return (ty,tstate3)
                  {-TODO handle what needs to happen with the wildcard argument, list of offsets in case of soa -}
                  AfterConstantLE i l1 ->
                       do r <- getRegion exp constrs l1
                          let tstate1 = extendTS loc (Output,True) $ setAfter l1 tstatein
                          let constrs1 = extendConstrs (InRegionC loc r) $ extendConstrs (AfterConstantC i l1 loc) constrs
                          (ty,tstate2) <- tcExp ddfs env' funs constrs1 regs tstate1 e
                          tstate3 <- removeLoc exp tstate2 loc
                          return (ty,tstate3)
                  AssignLE l1 -> 
                      do r <- getRegion exp constrs l1 
                         let tstate1 = extendTS loc (Output, True) $ setAfter l1 tstatein
                         let constrs1 = extendConstrs (InRegionC loc r) $ extendConstrs (AssignC l1 loc) constrs
                         (ty, tstate2) <- tcExp ddfs env' funs constrs1 regs tstate1 e
                         tstate3 <- removeLoc exp tstate2 loc
                         return (ty, tstate3)
                  AfterVariableLE x l1 _ ->
                      do r <- getRegion exp constrs l1
                         (_xty,tstate1) <- tcExp ddfs env funs constrs regs tstatein $ VarE x
                         -- NOTE: We now allow aliases (offsets) from scalar vars too. So we can leave out this check
                         -- ensurePackedLoc exp xty l1
                         let tstate2 = extendTS loc (Output,True) $ setAfter l1 tstate1
                         let constrs1 = extendConstrs (InRegionC loc r) $ extendConstrs (AfterVariableC x l1 loc) constrs
                         (ty,tstate3) <- tcExp ddfs env' funs constrs1 regs tstate2 e
                         tstate4 <- removeLoc exp tstate3 loc
                         return (ty,tstate4)
                  FromEndLE _l1 ->
                      do -- TODO: This is the bare minimum which gets the examples typechecking again.
                         -- Need to figure out if we need to check more things here
                          r <- getRegion exp constrs _l1
                          let tstate1 = extendTS loc (Output,True) $ setAfter _l1 tstatein
                          let constrs1 = extendConstrs (InRegionC loc r) $ extendConstrs (FromEndC _l1) constrs 
                          (ty,tstate2) <- tcExp ddfs env' funs constrs1 regs tstate1 e
                          return (ty,tstate2)
                  FreeLE ->
                      do let constrs1 = extendConstrs (InRegionC loc globalReg) $ constrs
                         (ty,tstate1) <- tcExp ddfs env' funs constrs1 regs tstatein e
                         return (ty,tstate1)
  
                  InRegionLE{} -> throwError $ GenericTC ("InRegionLE not handled.")  exp
                  -- Gen SoA Loc constratins is not used in the IR atm. 
                  -- GenSoALoc dcloc field_locs -> throwError $ GenericTC ("GenSoALoc not handled.") exp
                  GetDataConLocSoA soa_loc -> do 
                                          r <- getRegion exp constrs soa_loc
                                          -- get region for the data constructor buffer
                                          let r' = case r of 
                                                      SoAR dreg _ -> dreg
                                                      _ -> error $ "L2.Typecheck.tcExp: GetDataConLocSoA: Expected SoAR region, got " ++ show r
  
                                          -- LetLocE dconLoc = getDataConLocSoA soa_loc
                                          -- dbgTraceIt "print reg: " dbgTraceIt (sdoc r') dbgTraceIt "end print reg.\n" 
                                          let tstate1 = extendTS loc (Output,True) $ tstatein
                                          let constrs1 = extendConstrs (InRegionC loc r') $ constrs
                                          (ty,tstate2) <- tcExp ddfs env' funs constrs1 regs tstate1 e
                                          tstate3 <- removeLoc exp tstate2 loc
                                          return (ty,tstate3)
                  GetFieldLocSoA key soa_loc -> do
                                                -- get the region of the SoA loc. 
                                                r <- getRegion exp constrs soa_loc
                                                -- get the region of the field location 
                                                let Just r' = case r of 
                                                               SoAR _dreg fieldRegions -> lookup key fieldRegions
                                                               _ -> error $ "L2.Typecheck.tcExp: GetFieldLocSoA: Expected SoAR region, got " ++ show r
                                                let tstate1 = extendTS loc (Output, True) $ tstatein 
                                                let constrs1 = extendConstrs (InRegionC loc r') $ constrs
                                                (ty, tstate2) <- tcExp ddfs env' funs constrs1 regs tstate1 e
                                                tstate3 <- removeLoc exp tstate2 loc
                                                return (ty, tstate3)
                                                --throwError $ GenericTC ("GetFieldLocSoA not handled.") exp  
                  -- AfterVectorLE dexp rstExprs soa_loc -> throwError $ GenericTC ("AfterVectorLE not handled.") exp
  
        Ext (StartOfPkdCursor cur) -> do
          case M.lookup (fromVarToFreeVarsTy cur) (vEnv env) of
            Just (PackedTy{}) -> pure (CursorTy, tstatein)
            ty -> throwError $ GenericTC ("Expected PackedTy, got " ++ sdoc ty)  exp
  
        Ext (TagCursor a _b) -> do
          case M.lookup (fromVarToFreeVarsTy a) (vEnv env) of
            Just (PackedTy{}) -> pure (CursorTy, tstatein)
            ty -> throwError $ GenericTC ("Expected PackedTy, got " ++ sdoc ty)  exp
  
  
        Ext (FromEndE{}) -> throwError $ GenericTC "FromEndE not handled" exp
        Ext (AddFixed{}) -> return (CursorTy,tstatein)
  
        Ext (RetE _ls v) -> do
  
                 -- skip returned locations for now
                 recur tstatein $ VarE v
  
        -- The IntTy is just a placeholder. BoundsCheck is a side-effect
        Ext (BoundsCheck{}) -> return (IntTy,tstatein)
  
        Ext (IndirectionE tycon _ (a,_) _ _) -> return (PackedTy tycon a, tstatein)
  
        Ext GetCilkWorkerNum -> return (IntTy, tstatein)
  
        Ext (LetAvail _ e) -> recur tstatein e
  
        Ext (AllocateTagHere{}) -> do
          -- (ty,tstate1) <- recur tstatein (VarE v)
          -- ensureEqualTy (VarE v) ty CursorTy
          return (ProdTy [], tstatein)
  
        Ext (AllocateScalarsHere{}) -> do
          -- (ty,tstate1) <- recur tstatein (VarE v)
          -- ensureEqualTy (VarE v) ty CursorTy
          return (ProdTy [], tstatein)
  
        Ext (SSPush{}) -> do
          -- (ty,tstate1) <- recur tstatein (VarE v)
          -- ensureEqualTy (VarE v) ty CursorTy
          return (ProdTy [], tstatein)
  
        Ext (SSPop{}) -> do
          -- (ty,tstate1) <- recur tstatein (VarE v)
          -- ensureEqualTy (VarE v) ty CursorTy
          return (ProdTy [], tstatein)
  
        Ext (LetRegE _ _ _) -> throwError $ GenericTC "LetRegE not handled" exp
  
        Ext (BoundsCheckVector _) -> throwError $ GenericTC "BoundsCheckVector not handled" exp
  
      where recur ts e = tcExp ddfs env funs constrs regs ts e
            checkListElemTy el_ty =
              if isValidListElemTy el_ty
              then pure ()
              else throwError $ GenericTC ("Gibbon-TODO: Lists of only scalars or flat products of scalars are allowed. Got" ++ sdoc el_ty) exp
  
  
  
  -- | Helper function to check case branches.
  tcCases :: DDefs Ty2 -> Env2 FreeVarsTy Ty2 -> FunDefs2
          -> ConstraintSet -> RegionSet -> LocationTypeState -> LocVar
          -> Region -> [(DataCon, [(Var,LocVar)], Exp)]
          -> TcM ([Ty2], LocationTypeState)
  tcCases ddfs env funs constrs regs tstatein lin reg ((dc, vs, e):cases) = do
  
    let argtys = zip vs $ lookupDataCon ddfs dc
        argtys' = L.map (\((v, locvar), arg2) -> ((fromVarToFreeVarsTy v, locvar), arg2)) argtys
        pairwise = zip argtys $ Nothing : (L.map Just argtys)
  
        -- Generate the new constraints to check this branch
        genConstrs (((_v1,l1),PackedTy _ _),Nothing) (lin,lst) =
            (l1,[AfterConstantC 1 lin l1, InRegionC l1 reg] ++ lst)
        genConstrs (((_v1,l1),PackedTy _ _),Just ((v2,l2),PackedTy _ _)) (_lin,lst) =
            (l1,[AfterVariableC v2 l2 l1, InRegionC l1 reg] ++ lst)
        genConstrs (((_v1,l1),PackedTy _ _),Just ((_v2,_l2),IntTy)) (lin,lst) =
          let sz = fromMaybe 1 (sizeOfTy IntTy)
          in (l1, [AfterConstantC sz lin l1, InRegionC l1 reg] ++ lst)
        genConstrs (((_,l1),_),_) (lin,lst) =
          (lin, (InRegionC l1 reg : lst))
  
        -- Generate the new location state map to check this branch
        genTS ((_v,l),PackedTy _ _) ts = extendTS l (Input,False) ts
        genTS _ ts = ts
        genEnv ((v,l),PackedTy dc _l') env = extendVEnvLocVar v (PackedTy dc l) env
        genEnv ((v,_l),ty) env = extendVEnvLocVar v ty env
  
        -- Remove the pattern-bound location variables from the location state map
        remTS ((_v,l),PackedTy _ _) ts = removeTS l ts
        remTS _ ts = ts
  
        -- Use these functions with our old friend foldr
        constrs1 = L.foldr extendConstrs constrs $ snd $ L.foldr genConstrs (lin,[]) pairwise
        tstate1 = L.foldr genTS tstatein argtys
        env1 = L.foldr genEnv env argtys'
  
    (ty1,tstate2) <- tcExp ddfs env1 funs constrs1 regs tstate1 e
    (tyRest,tstateRest) <- recur
    tstateCombine <- combineTStates e tstate2 tstateRest
    let tstatee' = L.foldr remTS tstateCombine argtys'
    return (ty1:tyRest,tstatee')
  
      where recur = do
              (tys,tstate2) <- tcCases ddfs env funs constrs regs tstatein lin reg cases
              return (tys,tstate2)
  
  tcCases _ _ _ _ _ ts _ _ [] = return ([],ts)
  
  tcProj :: Exp -> Int -> Ty2 -> TcM Ty2
  tcProj _ i (ProdTy tys) = return $ tys !! i
  tcProj e _i ty = throwError $ GenericTC ("Projection from non-tuple type " ++ (show ty)) e
  
  -- | A wrapper for tcExp to check a list of exprs, checking them in order:
  -- the order matters because the location state map is threaded through,
  -- so this is assuming the list of expressions would have been evaluated
  -- in first-to-last order.
  tcExps :: DDefs Ty2 -> Env2 FreeVarsTy Ty2 -> FunDefs2
        -> ConstraintSet -> RegionSet -> LocationTypeState -> [Exp]
        -> TcM ([Ty2], LocationTypeState)
  tcExps ddfs env funs constrs regs tstatein (exp:exps) =
      do (ty,ts) <- tcExp ddfs env funs constrs regs tstatein exp
         (tys,ts') <- tcExps ddfs env funs constrs regs ts exps
         return (ty:tys,ts')
  tcExps _ _ _ _ _ ts [] = return ([],ts)
  
  
  
  -- | Main entry point, checks a whole program (functions and main body).
  tcProg :: Prog2 -> PassM Prog2
  tcProg prg0@Prog{ddefs,fundefs,mainExp} = do
  
    -- Handle functions
    mapM_ fd $ M.elems fundefs
  
    -- Handle main function
    case mainExp of
      Nothing -> return ()
      Just (e,t) ->
          let init_env = progToEnv' prg0
              res = runExcept $ tcExp ddefs init_env fundefs
                      (ConstraintSet $ S.empty) (RegionSet $ S.empty)
                      (LocationTypeState $ M.empty) e
          in case res of
               Left err -> error $ show err
               Right (t',_ts) ->
                   if t' == t then return ()
                   else error $ "Expected type " ++ (show t) ++ " and got type " ++ (show t')
  
    return prg0 -- Identity function for now.
  
    where
  
      fd :: FunDef2 -> PassM ()
      fd func@FunDef{funTy,funArgs,funBody} = do
          let init_env = progToEnv' prg0
              env = extendsVEnvLocVar (M.fromList $ zip (L.map fromVarToFreeVarsTy funArgs) (arrIns funTy)) init_env
              constrs = funConstrs (locVars funTy)
              regs = funRegs (locVars funTy)
              tstate = funTState (locVars funTy)
              res = runExcept $ tcExp ddefs env fundefs constrs regs tstate funBody
          case res of
            Left err -> error $ show err
            Right (ty,_) -> if ty == (arrOut funTy)
                            then return ()
                            else error $ "Expected type " ++ (sdoc (arrOut funTy))
                                      ++ " and got type " ++ (sdoc ty)
                                      ++ " in\n" ++ (sdoc func)
  
  
  
  --------------------------------------------------------------------------------------------
  
  -- Helper functions
  
  
  -- | Insert a region into a region set.
  -- Includes an expression for error reporting.
  regionInsert :: Exp -> Region -> RegionSet -> TcM RegionSet
  regionInsert e r (RegionSet regSet) = do
    if (S.member r regSet)
    then throwError $ GenericTC "Shadowed regions not allowed" e
    else return $ RegionSet (S.insert r regSet)
  
  -- | Ask if a region is in the region set.
  hasRegion :: Region -> RegionSet -> Bool
  hasRegion r (RegionSet regSet) = S.member r regSet
  
  -- | Ensure that a region is in a region set, reporting an error otherwise.
  -- Includes an expression for error reporting.
  ensureRegion :: Exp -> Region -> RegionSet -> TcM ()
  ensureRegion exp r (RegionSet regSet) =
      if S.member r regSet then return ()
      else throwError $ GenericTC ("Region " ++ (show r) ++ " not in scope") exp
  
  -- | Get the region of a location variable.
  -- Includes an expression for error reporting.
  getRegion :: Exp -> ConstraintSet -> LocVar -> TcM Region
  getRegion exp (ConstraintSet cs) l = go $ S.toList cs
      where go ((InRegionC l1 r):cs) = if l1 == l then return r
                                       else go cs
            go (_:cs) = go cs
            go [] = throwError $ GenericTC ("Location " ++ (show l) ++ " has no region") exp
  
  -- | Get the regions mentioned in the location bindings in a function type.
  funRegs :: [LRM] -> RegionSet
  funRegs ((LRM _l r _m):lrms) =
      let
        (RegionSet rs) = funRegs lrms
      in
        RegionSet $ S.insert r rs
  funRegs [] = RegionSet $ S.empty
  
  globalReg :: Region
  globalReg = GlobR "GLOBAL" BigInfinite
  
  -- | Get the constraints from the location bindings in a function type.
  funConstrs :: [LRM] -> ConstraintSet
  funConstrs ((LRM l r _m):lrms) = case r of 
                                      SoAR dclreg fieldRegs -> case l of 
                                                                  Single _ -> error "funConstrs: expected SoA location in LRM."
                                                                  SoA dloc flocs -> let nconstrs = extendConstrs (InRegionC (singleLocVar dloc) dclreg) $ funConstrs lrms
                                                                                        nconstrs' = extendConstrs (InRegionC l r) $ nconstrs
                                                                                       in L.foldr (\(k, freg) acc -> case lookup k flocs of
                                                                                                                          Nothing -> error "funConstrs: expected a field location for corresponding key in SoAR."
                                                                                                                          Just loc -> extendConstrs (InRegionC (loc) freg) acc
                                                                                                  ) nconstrs' fieldRegs
                                      _ -> extendConstrs (InRegionC l r) $ funConstrs lrms
  funConstrs [] = ConstraintSet $ S.empty
  
  -- | Get the type state implied by the location bindings in a function type.
  funTState :: [LRM] -> LocationTypeState
  funTState ((LRM l _r m):lrms) =
      extendTS l (m,False) $ funTState lrms
  funTState [] = LocationTypeState $ M.empty
  
  -- | Look up the type of a variable from the environment
  -- Includes an expression for error reporting.
  lookupVar :: Env2 FreeVarsTy Ty2 -> FreeVarsTy -> Exp -> TcM Ty2
  lookupVar env l _exp =
      case M.lookup l $ vEnv env of
        Nothing -> error $ "L2/Typecheck.hs: lookupVar: Variable not found in env."++ sdoc l
        Just ty -> return ty
  
  -- | Combine two location state maps.
  -- Currently just uses a naive union, which will not verify that different branches do the same
  -- things.
  -- TODO: Fix this!
  combineTStates :: Exp -> LocationTypeState -> LocationTypeState -> TcM LocationTypeState
  combineTStates _exp (LocationTypeState ts1) (LocationTypeState ts2) =
      return $ LocationTypeState $ M.union ts1 ts2
      -- throwError $ DivergingEffectsTC exp ts1 ts2
  
  -- | Ensure that two things are equal.
  -- Includes an expression for error reporting.
  ensureEqual :: Eq a => Exp -> String -> a -> a -> TcM a
  ensureEqual exp str a b = if a == b then return a else throwError $ GenericTC str exp
  
  -- | Ensure that the number of arguments to an operation is correct.
  checkLen :: (Show op, Show arg) => Exp -> op -> Int -> [arg] -> TcM ()
  checkLen expr pr n ls =
    if length ls == n
    then return ()
    else throwError $ GenericTC ("Wrong number of arguments to "++show pr++
                                 ".\nExpected "++show n++", received "
                                  ++show (length ls)++":\n  "++show ls) expr
  
  -- | Ensure that two types are equal.
  -- Includes an expression for error reporting.
  ensureEqualTy :: Exp -> Ty2 -> Ty2 -> TcM Ty2
  ensureEqualTy _ CursorTy IntTy = return CursorTy
  ensureEqualTy _ IntTy CursorTy = return CursorTy
  ensureEqualTy exp a b = ensureEqual exp ("Expected these types to be the same: "
                                           ++ (show a) ++ ", " ++ (show b)) a b
  
  ensureEqualTyModCursor :: Exp -> Ty2 -> Ty2 -> TcM Ty2
  ensureEqualTyModCursor _exp CursorTy (PackedTy _ _) = return CursorTy
  ensureEqualTyModCursor _exp (PackedTy _ _) CursorTy = return CursorTy
  ensureEqualTyModCursor exp a b = ensureEqualTy exp a b
  
  -- | Ensure that two types are equal, ignoring the locations if they are packed.
  -- Includes an expression for error reporting.
  ensureEqualTyNoLoc :: Exp -> Ty2 -> Ty2 -> TcM Ty2
  ensureEqualTyNoLoc exp ty1 ty2 =
    case (ty1,ty2) of
      (SymDictTy _ar1 ty1, SymDictTy _ar2 ty2) ->
          do ty1' <- dummyTyLocs ty1
             ty2' <- dummyTyLocs ty2
             ensureEqualTyNoLoc exp ty1' ty2'
      (PackedTy dc1 _, PackedTy dc2 _) -> if dc1 == dc2
                                          then return ty1
                                          else ensureEqualTy exp ty1 ty2
      (ProdTy tys1, ProdTy tys2) -> do
          checks <- return $ L.map (\(ty1,ty2) -> ensureEqualTyNoLoc exp ty1 ty2) (zip tys1 tys2)
          -- TODO: avoid runExcept here
          forM_ checks $ \c -> do
              let c' = runExcept c
              case c' of
                Left err -> throwError err
                Right _  -> return ()
          return ty1
      _ -> ensureEqualTyModCursor exp ty1 ty2
  
  
  -- | Ensure that match cases make sense.
  -- Includes an expression for error reporting.
  ensureMatchCases :: DDefs Ty2 -> Exp -> Ty2 -> [(DataCon, [(Var,LocVar)], Exp)] -> TcM ()
  ensureMatchCases ddfs exp ty cs = do
    case ty of
      PackedTy tc _l -> do
              let cons = S.fromList $ L.map fst $ dataCons $ lookupDDef ddfs tc
              forM cs $ \(dc,_,_) ->
                  do if S.member dc cons
                     then return ()
                     else throwError $ GenericTC "Invalid case statement" exp
              return ()
      _ -> throwError $ GenericTC "Cannot case on non-packed type" exp
  
  -- | Ensure a type is at a particular location.
  -- Includes an expression for error reporting.
  ensurePackedLoc :: Exp -> Ty2 -> LocVar -> TcM ()
  ensurePackedLoc exp ty l =
      case ty of
        PackedTy _ l1 -> if l1 == l then return ()
                         else throwError $ GenericTC ("Wrong location in type " ++ (show ty)) exp
        _ -> throwError $ GenericTC "Expected a packed type" exp
  
  -- need to check chain of aliased locations
  -- linit is the location in the type i.e lout653
  -- ensureAfterConstant expects l664 to be after lout653
  
  -- | Ensure the locations all line up with the constraints in a data constructor application.
  -- Includes an expression for error reporting.
  ensureDataCon :: Exp -> TyCon -> DataCon -> LocVar -> [Ty2] -> ConstraintSet -> TcM ()
  ensureDataCon exp dcty dc linit0 tys cs = case linit0 of
                                         Single _location -> (go Nothing linit0 tys)
                                            where 
                                              go Nothing linit ((PackedTy dc2 l):tys) = do
                                                ensureAfterConstant exp cs linit l
                                                go (Just (PackedTy dc2 l)) l tys
  
                                              go Nothing linit (_ty:tys) = do
                                                  case getAfterConstant cs linit of
                                                        Nothing     -> go Nothing linit tys
                                                        Just linit' -> go Nothing linit' tys
  
                                              go (Just (PackedTy _dc1 l1)) _linit ((PackedTy dc2 l2):tys) = do
                                                    ensureAfterPacked exp cs l1 l2
                                                    go (Just (PackedTy dc2 l2)) l2 tys
  
                                              go (Just (PackedTy _dc _l1)) linit (_ty:tys) =
                                               go Nothing linit tys
                                              go _ _ [] = return ()
                                              go _ _ _  = internalError "Unxpected case reached: L2:ensureDataCon"
  
                                         -- This checking should be fine for a Flat list data type
                                         -- data List = Cons Int List | Nil
                                         -- TODO: Extend for a Tree data type
                                         SoA dcloc fieldLocs -> do 
                                                let unself_idxs = L.concatMap 
                                                                        (\ty -> case ty of 
                                                                                    PackedTy k _ -> if k == dcty 
                                                                                                    then []
                                                                                                    else [fromJust (L.elemIndex ty tys)]
                                                                                    _ -> [fromJust (L.elemIndex ty tys)]
                                                  
                                                                        ) tys  
                                                let self_idxs = L.concatMap 
                                                                        (\ty -> case ty of 
                                                                                    PackedTy k _ -> if k == dcty 
                                                                                                    then [fromJust (L.elemIndex ty tys)]
                                                                                                    else []
                                                                                    _ -> []
                                                  
                                                                        ) tys  
                                                -- Not self recursive fields
                                                let _unselfTys = L.foldl (\a idx-> a ++ [tys !! idx]
                                                                                      ) [] unself_idxs
                                                -- Self recursive fields
                                                let selfTys = L.foldl (\a idx -> a ++ [tys !! idx]
                                                                                      ) [] self_idxs
  
                                                let _unselfWriteAtLocs = L.map (\idx -> lookup (dc, idx) fieldLocs) unself_idxs
                                                -- ensure after Constant with head of selfTys
                                                _ <- do 
                                                    case selfTys of 
                                                      [] -> return ()
                                                      x:_ -> case x of 
                                                               PackedTy _ l -> ensureAfterConstant exp cs (Single dcloc) (getDconLoc l) 
                                                               _ -> error "Did not expected unpacked type!"
                                                -- TODO: ensure after constant for all scalar not self recursive fields, with offset 0
                                                -- TODO: ensure after constant for all locs in dest with the next self recursive field. 
                                                _ <- do 
                                                    case selfTys of 
                                                      [] -> return () 
                                                      x:_ -> case x of
                                                                 PackedTy _ (SoA _dcloc' fieldLocs') -> do
                                                                                                          --let nextWriteAtLocs = L.map (\idx -> lookup (dc, idx) fieldLocs') unself_idxs
                                                                                                          let _aliasLocs = L.map (\idx -> case (lookup (dc, idx) fieldLocs) of
                                                                                                                                                            Just l -> getAfterConstantAlias cs l
                                                                                                                                                            Nothing -> error "enusreDataCon: Did not expect Nothing!"
                                                                                                                                      ) unself_idxs
                                                                                                          let _nextWriteAtLocs = L.map (\idx -> lookup (dc, idx) fieldLocs') unself_idxs
                                                                                                          -- dbgTraceIt "Print line 1241: " dbgTraceIt (sdoc (aliasLocs)) dbgTraceIt "End\n"
                                                                                                          {- VS typechecking fails? Why is this not working ??? -}
                                                                                                          -- _ <- mapM (\(Just l1, Just l2, ty) -> case ty of 
                                                                                                          --                                           PackedTy{} -> ensureAfterConstant exp cs l1 l2
                                                                                                          --                                           _ -> ensureAfterConstant exp cs l1 l2) (zip3 unselfWriteAtLocs aliasLocs unselfTys)
                                                                                                          -- -- dbgTraceIt "Print line 1241: " dbgTraceIt (sdoc (nextWriteAtLocs)) dbgTraceIt "End\n"
                                                                                                          -- _ <- mapM (\(Just l1, Just l2, ty) -> case ty of 
                                                                                                          --                                           PackedTy{} -> ensureAfterPacked exp cs l1 l2
                                                                                                          --                                           _ -> ensureAfterConstant exp cs l1 l2) (zip3 aliasLocs nextWriteAtLocs unselfTys)
                                                                                                          return ()
                                                                 _ -> error "ensureDataCon: Did not expected unpacked type!"
                                                -- dbgTraceIt "Print in ensure data con" dbgTraceIt (sdoc (unselfTys, selfTys, unselfWriteAtLocs)) dbgTraceIt "End\n"
                                                return ()
  
  
  -- | Ensure that one location is +c after another location in the constraint set.
  -- Includes an expression for error reporting.
  ensureAfterConstant :: Exp -> ConstraintSet -> LocVar -> LocVar -> TcM ()
  ensureAfterConstant exp (ConstraintSet cs) l1 l2 =
      if L.any f $ S.toList cs then return ()
      else throwError $ LocationTC "Expected after constant relationship" exp l1 l2
      -- l1 is before l2
      where f (AfterConstantC _i l1' l2') = l1' == l1 && l2' == l2
            f (AssignC l1' l2') = l1' == l1 && l2' == l2
            f _ = False
  
  -- | Ensure that one location is a variable size after another location in the constraint set.
  -- Includes an expression for error reporting.
  ensureAfterPacked :: Exp -> ConstraintSet -> LocVar -> LocVar -> TcM ()
  ensureAfterPacked  exp (ConstraintSet cs) l1 l2 =
      if L.any f $ S.toList cs then return ()
      else throwError $ LocationTC "Expected after packed relationship" exp l1 l2
      where f (AfterVariableC _v l1' l2') = l1' == l1 && l2' == l2
            f _ = False
  
  getAfterConstant :: ConstraintSet -> LocVar -> Maybe LocVar
  getAfterConstant (ConstraintSet cs) l0 =
    let mb_cs = L.find (\c -> case c of
                           AfterConstantC _i l1 _l2 | l1 == l0 -> True
                           _ -> False)
                cs
    in case mb_cs of
         Just (AfterConstantC _i _l1 l2) -> Just l2
         _ -> Nothing
  
  
  -- dbgTraceIt "Print in after alias:" dbgTraceIt (sdoc (c)) dbgTraceIt "End\n"
  getAfterConstantAlias :: ConstraintSet -> LocVar -> Maybe LocVar 
  getAfterConstantAlias (ConstraintSet cs) l0 = 
    let mb_cs = L.find (\c -> case c of
                           AfterConstantC i l1 _l2 | (l1 == l0 && i == 0) -> True
                           AssignC l1 _l2 | l1 == l0 -> True
                           _ -> False)
                cs
     in case mb_cs of 
          Just (AfterConstantC _i _l1 l2) -> Just l2
          Just (AssignC _l1 l2) -> Just l2
          _ -> Nothing  
  
  
  extendTS
    :: LocVar
       -> (Modality, Aliased) -> LocationTypeState -> LocationTypeState
  extendTS v d (LocationTypeState ls) = LocationTypeState $ M.insert v d ls
  
  removeTS :: LocVar -> LocationTypeState -> LocationTypeState
  removeTS l (LocationTypeState ls) = LocationTypeState $ M.delete l ls
  
  setAfter :: LocVar -> LocationTypeState -> LocationTypeState
  setAfter l (LocationTypeState ls) = LocationTypeState $ M.adjust (\(m,_) -> (m,True)) l ls
  
  _lookupTS :: Exp -> LocVar -> LocationTypeState -> TcM (Modality,Bool)
  _lookupTS exp l (LocationTypeState ls) =
      case M.lookup l ls of
        Nothing -> throwError $ GenericTC ("Failed lookup of location " ++ (show l)) exp
        Just d -> return d
  
  extendConstrs :: LocConstraint -> ConstraintSet -> ConstraintSet
  extendConstrs c (ConstraintSet cs) = ConstraintSet $ S.insert c cs
  
  switchOutLoc :: Exp -> LocationTypeState -> LocVar -> TcM LocationTypeState
  switchOutLoc exp ts@(LocationTypeState ls) l =
      case M.lookup l ls of
        Nothing -> throwError $ GenericTC ("Unknown location " ++ (show l)) exp
        Just (Output,a) -> return $ LocationTypeState $ M.update (\_ -> Just (Input,a)) l ls
        Just (Input,_a) -> return ts
  
  _absentAfter :: Exp -> LocationTypeState -> LocVar -> TcM ()
  _absentAfter exp (LocationTypeState ls) l =
      case M.lookup l ls of
        Nothing -> throwError $ GenericTC ("Unknown location " ++ (show l)) exp
        Just (_m,False) -> return ()
        Just (_m,True) -> throwError $ GenericTC ("Alias of location " ++ (show l)) exp
  
  -- | Ensure that a location is not already "defined" by a start constraint.
  absentStart :: Exp -> ConstraintSet -> Region -> TcM ()
  absentStart exp (ConstraintSet cs) r = go $ S.toList cs
      where go ((StartOfC _l r'):cs) =
                if r == r'
                then throwError $ GenericTC ("Repeated start of " ++ (show r)) exp
                else go cs
            go (_:cs) = go cs
            go [] = return ()
  
  
  removeLoc :: Exp -> LocationTypeState -> LocVar -> TcM LocationTypeState
  removeLoc exp (LocationTypeState ls) l =
      if M.member l ls
      then return $ LocationTypeState $ M.delete l ls
      else throwError $ GenericTC ("Cannot remove location " ++ (show l) ++ (show exp)) exp
  
  ensureArenaScope :: MonadError TCError m => Exp -> Env2 FreeVarsTy a -> Maybe Var -> m ()
  ensureArenaScope exp env ar =
      case ar of
        Nothing -> throwError $ GenericTC "Expected arena annotation" exp
        Just var -> unless (S.member (fromVarToFreeVarsTy var) . M.keysSet . vEnv $ env) $
                    throwError $ GenericTC ("Expected arena in scope: " ++ sdoc var) exp
  