content: {-# LANGUAGE CPP              #-}
  {-# LANGUAGE LambdaCase       #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE RecordWildCards  #-}
  
  module Gibbon.HaskellFrontend
    ( parseFile, primMap, multiArgsToOne, desugarLinearExts ) where
  
  import           Control.Monad
  import           Data.Foldable ( foldrM )
  #if !MIN_VERSION_base(4,21,0)
  import           Data.Foldable ( foldl' )
  #endif
  import           Data.Maybe (catMaybes, isJust)
  import qualified Data.Map as M
  import qualified Data.Set as S
  import qualified Safe as Sf
  
  import           Data.IORef
  import           Language.Haskell.Exts.Extension
  import           Language.Haskell.Exts.Parser
  import           Language.Haskell.Exts.Syntax as H
  import           Language.Haskell.Exts.Pretty
  import           Language.Haskell.Exts.SrcLoc
  import           Language.Haskell.Exts.CPP
  import           System.Environment ( getEnvironment )
  import           System.Directory
  import           System.FilePath
  import           System.Process
  import           System.Exit
  import           System.IO
  
  import           Gibbon.L0.Syntax as L0
  import           Gibbon.Common
  import           Gibbon.DynFlags
  
  --------------------------------------------------------------------------------
  
  {-
  
  Importing modules:
  ~~~~~~~~~~~~~~~~~~
  
  We use the same notion of search paths as GHC[1], except that GHC also has a
  set of "known" packages (base, containers, etc.) where it looks for modules.
  Gibbon doesn't have those, and the rootset for our search is a singleton {"."}.
  Consider this directory structure:
      .
      |── A
      |   └── B
      |       |── C.hs
      |       |── D.hs
      |       |── Foo.hs
      |── Bar.hs
  
  If Bar.hs has a `import A.B.C`, we look for a file `./A/B/C.hs`. However, note
  that this design is much more primitive than what Cabal/Stack allow. Can A.B.C
  import A.B.D? It depends on where we invoke GHC from. If we do it from ".", then
  yes, because A.B.D exists at A/B/D.hs. But if we run "ghc C.hs", it will fail since
  it expects A.B.D to be at A/B/A/B/D.hs.
  
  [1] https://downloads.haskell.org/ghc/8.6.4/docs/html/users_guide/separate_compilation.html?#the-search-path
  
  -}
  
  
  parseFile :: Config -> FilePath -> IO (PassM Prog0)
  parseFile cfg path = do
      pstate0_ref <- newIORef emptyParseState
      parseFile' cfg pstate0_ref [] path
  
  
  data ParseState = ParseState
      { imported :: M.Map (String, FilePath) Prog0 }
  
  emptyParseState :: ParseState
  emptyParseState = ParseState M.empty
  
  parseMode :: ParseMode
  parseMode = defaultParseMode { extensions = [ EnableExtension ScopedTypeVariables
                                              , EnableExtension CPP
                                              , EnableExtension TypeApplications
                                              ]
                                              ++ (extensions defaultParseMode)
                               }
  
  parseFile' :: Config -> IORef ParseState -> [String] -> FilePath -> IO (PassM Prog0)
  parseFile' cfg pstate_ref import_route path = do
    when (gopt Opt_GhcTc (dynflags cfg)) $
        typecheckWithGhc cfg path
    str <- readFile path
    let cleaned = removeLinearArrows str
    -- let parsed = parseModuleWithMode parseMode cleaned
    parsed <- parseFileContentsWithCommentsAndCPP defaultCpphsOptions parseMode cleaned
    case parsed of
      ParseOk (hs,_comments) -> desugarModule cfg pstate_ref import_route (takeDirectory path) hs
      ParseFailed loc er -> do
        error ("haskell-src-exts failed: " ++ er ++ ", at " ++ prettyPrint loc)
  
  -- | ASSUMPTION: gibbon-stdlib is available to Cabal.
  --
  -- Currently 'run_all_tests.sh' installs it with 'cabal v1-install . -w ghc-9.0.1'.
  typecheckWithGhc :: Config -> FilePath -> IO ()
  typecheckWithGhc cfg path = do
    when (verbosity cfg >= 3) $
      putStr " [compiler] Running pass, GHC typechecker\n   => "
    let cmd = "ghc-9.0.1 -package gibbon-stdlib-0.1 -XNoImplicitPrelude -fno-code " ++ path
    (_, Just hout, Just herr, phandle) <-
          createProcess (shell cmd)
              { std_out = CreatePipe
              , std_err = CreatePipe
              , cwd = Just (takeDirectory path)
              }
    exitCode <- waitForProcess phandle
    case exitCode of
      ExitSuccess -> do
        when (verbosity cfg >= 3) $ do
          out <- hGetContents hout
          err <- hGetContents herr
          putStrLn out
          putStrLn err
        pure ()
      ExitFailure _ -> do
        err <- hGetContents herr
        error err
  
  -- | Really basic, and won't catch every occurence of a linear arrow.
  --
  -- But its only a stop-gap until we move to ghc-lib-parser, which can parse
  -- linear types and other things not supported by haskell-src-exts (e.g. CPP).
  removeLinearArrows :: String -> String
  removeLinearArrows str =
      fst $
      foldr (\c (acc,saw_one) ->
                  if saw_one && c == '%'
                  then (acc, False)
                  else if saw_one && c /= '%'
                  then (c:'1':acc, False)
                  else if c == '1'
                  then (acc, True)
                  else (c:acc, False))
             ([],False)
             str
      {-
       - messup up indendataion and causes compilation errors.
       -
       - unlines .
       - map (unwords .
       -      map (\w -> if w == "%1->" || w == "%1 ->"
       -                 then "->"
       -                 else w) .
       -      words) .
       - lines
       -}
  
  data TopLevel
    = HDDef (DDef Ty0)
    | HFunDef (FunDef Var Exp0)
    | HMain (Maybe (Exp0, Ty0))
    | HInline Var
    | MemLayoutTy TyCon MemoryLayout 
    deriving (Show, Eq)
  
  type TopTyEnv = TyEnv Var TyScheme
  type TypeSynEnv = M.Map TyCon Ty0
  
  desugarModule :: (Show a,  Pretty a)
                => Config -> IORef ParseState -> [String] -> FilePath -> Module a -> IO (PassM Prog0)
  desugarModule cfg pstate_ref import_route dir (Module _ head_mb _pragmas imports decls) = do
    let type_syns = foldl collectTypeSynonyms M.empty decls
        -- Since top-level functions and their types can't be declared in
        -- single top-level declaration we first collect types and then collect
        -- definitions.
        funtys = foldr (collectTopTy type_syns) M.empty decls
    imported_progs :: [PassM Prog0] <- mapM (processImport cfg pstate_ref (mod_name : import_route) dir) imports
    let prog = do
          toplevels <- catMaybes <$> mapM (collectTopLevel type_syns funtys) decls
          let (defs,_vars,funs,inlines,main, memlayouts) = foldr classify init_acc toplevels
              defs' = updateMemoryLayout defs memlayouts
              funs' = foldr (\v acc -> M.update (\fn@(FunDef{funMeta}) -> Just (fn { funMeta = funMeta { funInline = Inline }})) v acc) funs inlines
          imported_progs' <- mapM id imported_progs
          let (defs0,funs0) =
                foldr
                  (\Prog{ddefs,fundefs} (defs1,funs1) ->
                       let ddef_names1 = M.keysSet defs1
                           ddef_names2 = M.keysSet ddefs
                           fn_names1 = M.keysSet funs1
                           fn_names2 = M.keysSet fundefs
                           em1 = S.intersection ddef_names1 ddef_names2
                           em2 = S.intersection fn_names1 fn_names2
                           conflicts1 = foldr
                                          (\d acc ->
                                               if (ddefs M.! d) /= (defs1 M.! d)
                                               then d : acc
                                               else acc)
                                          []
                                          em1
                           conflicts2 = foldr
                                          (\f acc ->
                                               if (fundefs M.! f) /= (funs1 M.! f)
                                               then dbgTraceIt (sdoc ((fundefs M.! f), (funs1 M.! f))) (f : acc)
                                               else acc)
                                          []
                                          em2
                       in case (conflicts1, conflicts2) of
                              ([], []) -> (M.union ddefs defs1,  M.union fundefs funs1)
                              (_x:_xs,_) -> error $ "Conflicting definitions of " ++ show conflicts1 ++ " found in " ++ mod_name
                              (_,_x:_xs) -> error $ "Conflicting definitions of " ++ show (S.toList em2) ++ " found in " ++ mod_name)
                  (defs', funs')
                  imported_progs'
          pure (Prog defs0 funs0 main)
    pure prog
    where
      init_acc = (M.empty, M.empty, M.empty, S.empty, Nothing, M.empty)
      mod_name = moduleName head_mb
  
      moduleName :: Maybe (ModuleHead a) -> String
      moduleName Nothing = "Main"
      moduleName (Just (ModuleHead _ mod_name1 _warnings _exports)) =
        mnameToStr mod_name1
  
      classify thing (defs,vars,funs,inlines,main, memlayouts) =
        case thing of
          HDDef d   -> (M.insert (tyName d) d defs, vars, funs, inlines, main, memlayouts)
          HFunDef f -> (defs, vars, M.insert (funName f) f funs, inlines, main, memlayouts)
          HMain m ->
            case main of
              Nothing -> (defs, vars, funs, inlines, m, memlayouts)
              Just _  -> error $ "A module cannot have two main expressions."
                                 ++ show mod_name
          HInline v   -> (defs,vars,funs,S.insert v inlines,main, memlayouts)
          MemLayoutTy tycon l -> (defs,vars,funs,inlines,main, M.insert tycon l memlayouts)
      
      updateMemoryLayout indefs memlayouts = 
        let defs'' = M.mapWithKey (\k v -> let tyName = fromVar k
                                               layout = M.lookup tyName memlayouts
                                            in case layout of 
                                                    Just val -> let 
                                                                  v' = v{memLayout=val}
                                                                 in v'  
                                                    Nothing -> let 
                                                                  v' = v{memLayout=Linear}
                                                                 in v'
                                 ) indefs
          in defs''
      
  desugarModule _ _ _ _ m = error $ "desugarModule: " ++ prettyPrint m
  
  stdlibModules :: [String]
  stdlibModules =
    [ "Gibbon.Prim"
    , "Gibbon.Prelude"
    , "Gibbon.Vector"
    , "Gibbon.Vector.Parallel"
    , "Gibbon.List"
    , "Gibbon.PList"
    , "Gibbon.ByteString"
    ]
  
  processImport :: Config -> IORef ParseState -> [String] -> FilePath -> ImportDecl a -> IO (PassM Prog0)
  processImport cfg pstate_ref import_route dir decl@ImportDecl{..}
      -- When compiling with Gibbon, we should *NOT* inline things defined in Gibbon.Prim.
      | mod_name == "Gibbon.Prim" = pure (pure (Prog M.empty M.empty Nothing))
      | otherwise = do
      when (mod_name `elem` import_route) $
        error $ "Circular dependency detected. Import path: "++ show (mod_name : import_route)
      when (importQualified) $ error $ "Qualified imports not supported yet. Offending import: " ++  prettyPrint decl
      when (isJust importAs) $ error $ "Module aliases not supported yet. Offending import: " ++  prettyPrint decl
      when (isJust importSpecs) $ error $ "Selective imports not supported yet. Offending import: " ++  prettyPrint decl
      (ParseState imported) <- readIORef pstate_ref
      mod_fp <- if mod_name `elem` stdlibModules
                  then stdlibImportPath mod_name
                  else modImportPath importModule dir
      dbgTrace 5 ("Looking at " ++ mod_name) (pure ())
      dbgTrace 5 ("Previously imported: " ++ show (M.keysSet imported)) (pure ())
      prog <- case M.lookup (mod_name, mod_fp) imported of
                  Just prog -> do
                      dbgTrace 5 ("Already imported " ++ mod_name) (pure ())
                      pure prog
                  Nothing -> do
                      dbgTrace 5 ("Importing " ++ mod_name ++ " from " ++ mod_fp) (pure ())
                      prog0 <- parseFile' cfg pstate_ref import_route mod_fp
                      (ParseState imported') <- readIORef pstate_ref
                      let (prog0',_) = defaultRunPassM prog0
                      let imported'' = M.insert (mod_name, mod_fp) prog0' imported'
                      let pstate' = ParseState { imported = imported'' }
                      writeIORef pstate_ref pstate'
                      pure prog0'
  
      pure (pure prog)
    where
      mod_name = mnameToStr importModule
  
  stdlibImportPath :: String -> IO FilePath
  stdlibImportPath mod_name = do
      env <- getEnvironment
      let stdlibPath = case lookup "GIBBONDIR" env of
                      Just p -> p </> "gibbon-stdlib" </> modNameToFilename mod_name
                      -- Assume we're running from the compiler dir!
                      Nothing -> modNameToFilename mod_name
      e <- doesFileExist stdlibPath
      unless e $ error$ "stdlib.hs file not found at path: "++stdlibPath
                       ++"\n Consider setting GIBBONDIR to repo root.\n"
      pure stdlibPath
    where
      modNameToFilename :: String -> String
      modNameToFilename "Gibbon.Prelude" = "Gibbon" </> "Prelude.hs"
      modNameToFilename "Gibbon.Vector" = "Gibbon" </> "Vector.hs"
      modNameToFilename "Gibbon.Vector.Parallel" = "Gibbon" </> "Vector" </> "Parallel.hs"
      modNameToFilename "Gibbon.List" = "Gibbon" </> "List.hs"
      modNameToFilename "Gibbon.PList" = "Gibbon" </> "PList.hs"
      modNameToFilename "Gibbon.ByteString" = "Gibbon" </> "ByteString.hs"
      modNameToFilename oth = error $ "Unknown module: " ++ oth
  
  modImportPath :: ModuleName a -> String -> IO FilePath
  modImportPath importModule dir = do
      let mod_name = mnameToStr importModule
      mb_fp <- findModule dir importModule
      case mb_fp of
          Nothing -> error $ "Cannot find module: " ++
                     show mod_name ++ " in " ++ dir
          Just mod_fp -> pure mod_fp
  
  -- | Look for a module on the filesystem.
  findModule :: FilePath -> ModuleName a -> IO (Maybe FilePath)
  findModule dir m = do
    let mod_fp  = dir </> moduleNameToSlashes m <.> "hs"
    doesFileExist mod_fp >>= \b ->
      if b
      then pure $ Just mod_fp
      else pure Nothing
  
  -- | Returns the string version of the module name, with dots replaced by slashes.
  --
  moduleNameToSlashes :: ModuleName a -> String
  moduleNameToSlashes (ModuleName _ s) = dots_to_slashes s
    where dots_to_slashes = map (\c -> if c == '.' then pathSeparator else c)
  
  
  builtinTys :: S.Set Var
  builtinTys = S.fromList $
      [ "Int", "Float", "Bool", "Sym", "SymHash", "IntHash", "SymSet", "SymDict", "Arena", "Vector" ]
  
  keywords :: S.Set Var
  keywords = S.fromList $ map toVar $
      -- These cannot be added to primMap because they all require special handling while parsing.
      --
      [ "quote", "bench", "error", "par", "spawn", "is_big"
      -- operations on vectors
      , "valloc", "vnth", "vlength", "vslice", "inplacevupdate",
        "vsort", "inplacevsort", "vfree", "vfree2"
      -- parallel dictionaries
      , "alloc_pdict", "insert_pdict", "lookup_pdict", "member_pdict", "fork_pdict", "join_pdict"
      -- linked lists
      , "alloc_ll", "is_empty_ll", "cons_ll", "head_ll", "tail_ll", "free_ll", "free2_ll", "copy_ll"
      ] ++ M.keys primMap
  
  desugarTopType :: (Show a,  Pretty a) => TypeSynEnv -> Type a -> TyScheme
  desugarTopType type_syns ty =
    case ty of
      -- forall tvs ty.
      TyForall _ mb_tvbind _ ty1 ->
        let tyvars = case mb_tvbind of
                       Just bnds -> map desugarTyVarBind bnds
                       Nothing   -> []
        in ForAll tyvars (desugarType type_syns ty1)
      -- quantify over all tyvars.
      _ -> let ty' = desugarType type_syns ty
               tyvars = tyVarsInTy ty'
          in ForAll tyvars ty'
  
  desugarType :: (Show a,  Pretty a) => TypeSynEnv -> Type a -> Ty0
  desugarType type_syns ty =
    case ty of
      H.TyVar _ (Ident _ t) -> L0.TyVar $ UserTv (toVar t)
      TyTuple _ Boxed tys   -> ProdTy (map (desugarType type_syns) tys)
      TyCon _ (Special _ (UnitCon _))     -> ProdTy []
      TyCon _ (UnQual _ (Ident _ "Int"))  -> IntTy
      TyCon _ (UnQual _ (Ident _ "Char")) -> CharTy
      TyCon _ (UnQual _ (Ident _ "Float"))-> FloatTy
      TyCon _ (UnQual _ (Ident _ "Bool")) -> BoolTy
      TyCon _ (UnQual _ (Ident _ "Sym"))  -> SymTy0
      TyCon _ (UnQual _ (Ident _ "SymSet"))  -> SymSetTy
      TyCon _ (UnQual _ (Ident _ "SymHash"))  -> SymHashTy
      TyCon _ (UnQual _ (Ident _ "IntHash"))  -> IntHashTy
      TyCon _ (UnQual _ (Ident _ con)) ->
        case M.lookup con type_syns of
          Nothing -> PackedTy con []
          Just ty' -> ty'
      TyFun _ t1 t2 -> let t1' = desugarType type_syns t1
                           t2' = desugarType type_syns t2
                       in ArrowTy [t1'] t2'
      TyParen _ ty1 -> desugarType type_syns ty1
      TyApp _ tycon arg ->
        let ty' = desugarType type_syns tycon in
        case ty' of
          PackedTy con tyargs ->
              case (con,tyargs) of
                  ("Vector",[]) -> VectorTy (desugarType type_syns arg)
                  ("List",[]) -> ListTy (desugarType type_syns arg)
                  ("PDict",[]) ->
                    let arg' = desugarType type_syns arg in
                    case arg' of
                      ProdTy [k, v] -> PDictTy k v
                      _ -> error $ "desugarType: Unexpected PDictTy argument: " ++ show arg'
                  _ ->
                    case M.lookup con type_syns of
                      Nothing -> PackedTy con (tyargs ++ [desugarType type_syns arg])
                      Just ty'' -> ty''
          _ -> error $ "desugarType: Unexpected type arguments: " ++ show ty'
      _ -> error $ "desugarType: Unsupported type: " ++ show ty
  
  
  -- Like 'desugarTopType' but understands boxity.
  desugarTopType' :: (Show a,  Pretty a) => TypeSynEnv -> Type a -> (IsBoxed, TyScheme)
  desugarTopType' type_syns ty =
    case ty of
      -- forall tvs ty.
      TyForall _ mb_tvbind _ ty1 ->
        let tyvars = case mb_tvbind of
                       Just bnds -> map desugarTyVarBind bnds
                       Nothing   -> []
            (boxity, ty') = desugarType' type_syns ty1
        in (boxity, ForAll tyvars ty')
      -- quantify over all tyvars.
      _ -> let (boxity, ty') = desugarType' type_syns ty
               tyvars = tyVarsInTy ty'
          in (boxity, ForAll tyvars ty')
  
  -- Like 'desugarType' but understands boxity.
  desugarType' :: (Show a,  Pretty a) => TypeSynEnv -> Type a -> (IsBoxed, Ty0)
  desugarType' type_syns ty =
    case ty of
      TyBang _ _ (NoUnpack _) ty1 -> (True, desugarType type_syns ty1)
      _ -> (False, desugarType type_syns ty)
  
  -- | Transform a multi-argument function type to one where all inputs are a
  -- single tuple argument. E.g. (a -> b -> c -> d) => ((a,b,c) -> d).
  unCurryTopTy :: TyScheme -> TyScheme
  unCurryTopTy (ForAll tyvars ty) = ForAll tyvars (unCurryTy ty)
  
  unCurryTy :: Ty0 -> Ty0
  unCurryTy ty1 =
    case ty1 of
      ArrowTy _ ArrowTy{} ->
        let (a,b) = go [] ty1
            a' = map unCurryTy a
        in ArrowTy a' b
      _ -> ty1
    where
      go :: [Ty0] -> Ty0 -> ([Ty0], Ty0)
      go acc ty =
        case ty of
          ArrowTy as b -> (go (acc++as) b)
          _ -> (acc,ty)
  
  -- ^ A map between SExp-frontend prefix function names, and Gibbon
  -- abstract Primops.
  primMap :: M.Map String (Prim a)
  primMap = M.fromList
    [ ("+", AddP)
    , ("-", SubP)
    , ("*", MulP)
    , ("/", DivP)
    , ("div", DivP)
    , ("^", ExpP)
    , (".+.", FAddP)
    , (".-.", FSubP)
    , (".*.", FMulP)
    , ("./.", FDivP)
    , ("sqrt", FSqrtP)
    , ("==", EqIntP)
    , (".==.", EqFloatP)
    , ("*==*", EqCharP)
    , ("<", LtP)
    , (">", GtP)
    , ("<=", LtEqP)
    , (">=", GtEqP)
    , (".<.", FLtP)
    , (".>.", FGtP)
    , (".<=.", FLtEqP)
    , (".>=.", FGtEqP)
    , ("tan", FTanP)
    , ("mod", ModP)
    , ("||" , OrP)
    , ("&&", AndP)
    , ("eqsym", EqSymP)
    , ("rand", RandP)
    , ("frand", FRandP)
    , ("intToFloat", IntToFloatP)
    , ("floatToInt", FloatToIntP)
    , ("sizeParam", SizeParam)
    , ("getNumProcessors", GetNumProcessors)
    , ("True", MkTrue)
    , ("False", MkFalse)
    , ("gensym", Gensym)
    , ("printint", PrintInt)
    , ("printchar", PrintChar)
    , ("printfloat", PrintFloat)
    , ("printbool", PrintBool)
    , ("printsym", PrintSym)
    , ("readint", ReadInt)
    , ("is_big", IsBig)
    , ("empty_set", SymSetEmpty)
    , ("insert_set", SymSetInsert)
    , ("contains_set", SymSetContains)
    , ("empty_hash", SymHashEmpty)
    , ("insert_hash", SymHashInsert)
    , ("lookup_hash", SymHashLookup)
    , ("contains_hash", SymHashContains)
    , ("empty_int_hash", IntHashEmpty)
    , ("insert_int_hash", IntHashInsert)
    , ("lookup_int_hash", IntHashLookup)
    ]
  
  desugarExp :: (Show a, Pretty a) => TypeSynEnv -> TopTyEnv -> Exp a -> PassM Exp0
  desugarExp type_syns toplevel e =
    case e of
      Paren _ (ExpTypeSig _ (App _ (H.Var _ f) (Lit _ lit)) tyc)
          | (qnameToStr f) == "error" -> pure $ PrimAppE (ErrorP (litToString lit) (desugarType type_syns tyc)) []
      -- Paren _ (App _ (H.Var _ f) (Lit _ lit))
      --     | (qnameToStr f) == "error" -> pure $ PrimAppE (ErrorP (litToString lit
      Paren _ e2 -> desugarExp type_syns toplevel e2
      H.Var _ qv -> do
        let str = qnameToStr qv
            v = (toVar str)
        if str == "alloc_pdict"
        then do
          kty  <- newMetaTy
          vty  <- newMetaTy
          pure $ PrimAppE (PDictAllocP kty vty) []
        else if str == "alloc_ll"
        then do
          ty  <- newMetaTy
          pure $ PrimAppE (LLAllocP ty) []
        else if v == "sync"
        then pure SyncE
        else if v == "lsync"
        then pure SyncE
        else if M.member str primMap
        then pure $ PrimAppE (primMap M.! str) []
        else case M.lookup v toplevel of
               Just sigma ->
                 case tyFromScheme sigma of
                   ArrowTy{} ->
                     -- Functions with >0 args must be VarE's here -- the 'App _ e1 e2'
                     -- case below depends on it.
                     pure $ VarE v
                   -- Otherwise, 'v' is a top-level value binding, which we
                   -- encode as a function which takes no arguments.
                   _ -> pure $ AppE v [] []
               Nothing -> pure $ VarE v
      Lit _ lit  -> desugarLiteral lit
  
      Lambda _ pats bod -> do
        bod' <- desugarExp type_syns toplevel bod
        (vars,tys,bindss) <- unzip3 <$> mapM (desugarPatWithTy type_syns) pats
        let binds = concat bindss
            args = zip vars tys
        pure $ Ext $ LambdaE args (mkLets binds bod')
  
      App _ e1 e2 -> do
          desugarExp type_syns toplevel e1 >>= \case
            (VarE f) ->
              case M.lookup (fromVar f) primMap of
                Just p  -> (\e2' -> PrimAppE p [e2']) <$> desugarExp type_syns toplevel e2
                Nothing ->
                    if f == "quote"
                    then case e2 of
                           Lit _ lit -> pure $ LitSymE (toVar $ litToString lit)
                           _ -> error "desugarExp: quote only accepts string literals. E.g quote \"hello\""
                    else if f == "eqBenchProg"
                    then case e2 of
                           Lit _ lit -> pure $ (PrimAppE (EqBenchProgP (litToString lit)) [])
                           _ -> error "desugarExp: eqBenchProg only accepts string literals."
                    else if f == "readArrayFile"
                    then let go e0 = case e0 of
                                      Con _ (UnQual _ (Ident _ "Nothing")) -> do
                                        t <- newMetaTy
                                        pure $ PrimAppE (ReadArrayFile Nothing t) []
                                      App _ (Con _ (UnQual _ (Ident _ "Just"))) (Tuple _ Boxed [Lit _ name, Lit _ len]) -> do
                                        t <- newMetaTy
                                        pure $ PrimAppE (ReadArrayFile (Just (litToString name, litToInt len)) t) []
                                      Paren _ e3 -> go e3
                                      _ -> error $ "desugarExp: couldn't parse readArrayFile; " ++ show e0
                         in go e2
                    else if f == "readPackedFile"
                    then let go e0 = case e0 of
                                       TypeApp _ (TyCon _ (UnQual _ (Ident _ con))) -> do
                                         let ty = PackedTy con []
                                         pure $ PrimAppE (ReadPackedFile Nothing con Nothing ty) []
                                       _ -> error $ "desugarExp: couldn't parse readPackedFile; " ++ show e0
                         in go e2
                    else if f == "writePackedFile"
                    then
                      case e2 of
                        Lit _ fp -> do
                          ty <- newMetaTy
                          pure $ PrimAppE (WritePackedFile (litToString fp) ty) []
                        _ -> error $ "desugarExp: couldn't parse writePackedFile; " ++ show e2
                    else if f == "bench"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      pure $ Ext $ BenchE "HOLE" [] [e2'] False
                    else if f == "timeit"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty <- newMetaTy
                      pure $ TimeIt e2' ty False
                    else if f == "iterate"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty <- newMetaTy
                      pure $ TimeIt e2' ty True
                    else if f == "error"
                    then case e2 of
                           Lit _ lit -> pure $ PrimAppE (ErrorP (litToString lit) IntTy) [] -- assume int (!)
                           _ -> error "desugarExp: error expects String literal."
                    else if f == "par"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      pure $ Ext $ ParE0 [e2']
                    else if f == "spawn"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      pure $ SpawnE "HOLE" [] [e2']
                    else if f == "valloc"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty  <- newMetaTy
                      pure $ PrimAppE (VAllocP ty) [e2']
                    else if f == "vfree"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty  <- newMetaTy
                      pure $ PrimAppE (VFreeP ty) [e2']
                    else if f == "vfree2"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty  <- newMetaTy
                      pure $ PrimAppE (VFree2P ty) [e2']
                    else if f == "vnth"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty  <- newMetaTy
                      pure $ PrimAppE (VNthP ty) [e2']
                    else if f == "vlength"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty  <- newMetaTy
                      pure $ PrimAppE (VLengthP ty) [e2']
                    else if f == "inplacevupdate"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty  <- newMetaTy
                      pure $ PrimAppE (InplaceVUpdateP ty) [e2']
                    else if f == "vconcat"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty  <- newMetaTy
                      pure $ PrimAppE (VConcatP ty) [e2']
                    else if f == "vsort"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty  <- newMetaTy
                      pure $ PrimAppE (VSortP ty) [e2']
                    else if f == "inplacevsort"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty  <- newMetaTy
                      pure $ PrimAppE (InplaceVSortP ty) [e2']
                    else if f == "vslice"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty  <- newMetaTy
                      pure $ PrimAppE (VSliceP ty) [e2']
  
                    else if f == "vmerge"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty  <- newMetaTy
                      pure $ PrimAppE (VMergeP ty) [e2']
                    else if f == "insert_pdict"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      kty  <- newMetaTy
                      vty  <- newMetaTy
                      pure $ PrimAppE (PDictInsertP kty vty) [e2']
  
                    else if f == "lookup_pdict"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      kty  <- newMetaTy
                      vty  <- newMetaTy
                      pure $ PrimAppE (PDictLookupP kty vty) [e2']
                    else if f == "member_pdict"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      kty  <- newMetaTy
                      vty  <- newMetaTy
                      pure $ PrimAppE (PDictHasKeyP kty vty) [e2']
                    else if f == "fork_pdict"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      kty  <- newMetaTy
                      vty  <- newMetaTy
                      pure $ PrimAppE (PDictForkP kty vty) [e2']
                    else if f == "join_pdict"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      kty  <- newMetaTy
                      vty  <- newMetaTy
                      pure $ PrimAppE (PDictJoinP kty vty) [e2']
                    else if f == "is_empty_ll"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty  <- newMetaTy
                      pure $ PrimAppE (LLIsEmptyP ty) [e2']
                    else if f == "cons_ll"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty  <- newMetaTy
                      pure $ PrimAppE (LLConsP ty) [e2']
                    else if f == "head_ll"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty  <- newMetaTy
                      pure $ PrimAppE (LLHeadP ty) [e2']
                    else if f == "tail_ll"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty  <- newMetaTy
                      pure $ PrimAppE (LLTailP ty) [e2']
                    else if f == "free_ll"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty  <- newMetaTy
                      pure $ PrimAppE (LLFreeP ty) [e2']
                    else if f == "free2_ll"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty  <- newMetaTy
                      pure $ PrimAppE (LLFree2P ty) [e2']
                    else if f == "copy_ll"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty  <- newMetaTy
                      pure $ PrimAppE (LLCopyP ty) [e2']
                    else if f == "fst"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      pure $ ProjE 0 e2'
                    else if f == "snd"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      pure $ ProjE 1 e2'
                    else if f == "printPacked"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty <- newMetaTy
                      pure $ Ext (PrintPacked ty e2')
                    else if f == "copyPacked"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty <- newMetaTy
                      pure $ Ext (CopyPacked ty e2')
                    else if f == "travPacked"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      ty <- newMetaTy
                      pure $ Ext (TravPacked ty e2')
                    else if f == "unsafeAlias"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      pure $ Ext (LinearExt (AliasE e2'))
                    else if f == "unsafeToLinear"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      pure $ Ext (LinearExt (ToLinearE e2'))
                    else if f == "lseq"
                    then do
                      e2' <- desugarExp type_syns toplevel e2
                      pure $ Ext (LinearExt (LseqE e2' undefined))
                    else if S.member f keywords
                    then error $ "desugarExp: Keyword not handled: " ++ sdoc f
                    else AppE f [] <$> (: []) <$> desugarExp type_syns toplevel e2
            (DataConE tyapp c as) -> (\e2' -> DataConE tyapp c (as ++ [e2'])) <$> desugarExp type_syns toplevel e2
            (Ext (ParE0 ls)) -> do
              e2' <- desugarExp type_syns toplevel e2
              pure $ Ext $ ParE0 (ls ++ [e2'])
            (AppE f [] ls) -> do
              e2' <- desugarExp type_syns toplevel e2
              pure $ AppE f [] (ls ++ [e2'])
  
            (Ext (BenchE fn [] ls b)) -> do
              e2' <- desugarExp type_syns toplevel e2
              pure $ Ext $ BenchE fn [] (ls ++ [e2']) b
  
            (SpawnE fn [] ls) -> do
              e2' <- desugarExp type_syns toplevel e2
              pure $ SpawnE fn [] (ls ++ [e2'])
  
            (PrimAppE (WritePackedFile fp ty) ls) -> do
               e2' <- desugarExp type_syns toplevel e2
               pure $ PrimAppE (WritePackedFile fp ty) (ls ++ [e2'])
  
            (PrimAppE (ReadPackedFile _mb_fp tycon mb_var ty) []) ->
               let go e0 = case e0 of
                             Con _ (UnQual _ (Ident _ "Nothing")) -> do
                               pure (PrimAppE (ReadPackedFile Nothing tycon mb_var ty) [])
                             App _ (Con _ (UnQual _ (Ident _ "Just"))) (Lit _ name) -> do
                               pure (PrimAppE (ReadPackedFile (Just (litToString name)) tycon mb_var ty) [])
                             Paren _ e3 -> go e3
                             _ -> error $ "desugarExp: couldn't parse readPackedFile; " ++ show e0
               in go e2
  
            (PrimAppE (VMergeP elty) ls) -> do
              e2' <- desugarExp type_syns toplevel e2
              pure $ PrimAppE (VMergeP elty) (ls ++ [e2'])
            (PrimAppE p ls) -> do
              e2' <- desugarExp type_syns toplevel e2
              pure $ PrimAppE p (ls ++ [e2'])
  
            TimeIt{} ->
              error "desugarExp: TimeIt can only accept 1 expression."
  
            (Ext (LinearExt (LseqE a _))) -> do
              e2' <- desugarExp type_syns toplevel e2
              pure (Ext (LinearExt (LseqE a e2')))
  
            (Ext (LinearExt (ToLinearE (AppE f [] ls)))) -> do
              e2' <- desugarExp type_syns toplevel e2
              pure (Ext (LinearExt (ToLinearE (AppE f [] (ls ++ [e2'])))))
  
            (Ext (LinearExt (ToLinearE (DataConE tyapp dcon ls)))) -> do
              e2' <- desugarExp type_syns toplevel e2
              pure (Ext (LinearExt (ToLinearE (DataConE tyapp dcon (ls ++ [e2'])))))
  
            (Ext (LinearExt (ToLinearE (Ext (LambdaE [(v,ty)] bod))))) -> do
              e2' <- desugarExp type_syns toplevel e2
              pure (Ext (LinearExt (ToLinearE (LetE (v,[],ty,e2') bod))))
  
            (Ext (LinearExt (ToLinearE (VarE fn)))) -> do
              e2' <- desugarExp type_syns toplevel e2
              pure (Ext (LinearExt (ToLinearE (AppE fn [] [e2']))))
  
            f -> error ("desugarExp: Couldn't parse function application: (" ++ show f ++ ")")
  
      Let _ (BDecls _ decls) rhs -> do
        rhs' <- desugarExp type_syns toplevel rhs
        let funtys = foldr (collectTopTy type_syns) M.empty decls
        foldrM (generateBind type_syns toplevel funtys) rhs' decls
  
      If _ a b c -> do
        a' <- desugarExp type_syns toplevel a
        b' <- desugarExp type_syns toplevel b
        c' <- desugarExp type_syns toplevel c
        pure $ IfE a' b' c'
  
      Tuple _ Unboxed _ -> error $ "desugarExp: Only boxed tuples are allowed: " ++ prettyPrint e
      Tuple _ Boxed es  -> MkProdE <$> mapM (desugarExp type_syns toplevel) es
  
      Case _ scrt alts -> do
        scrt' <- desugarExp type_syns toplevel scrt
        CaseE scrt' <$> mapM (desugarAlt type_syns toplevel) alts
  
      Con _ (Special _ (UnitCon _)) -> pure $ MkProdE []
  
      Con _ qname -> do
        let dcon = qnameToStr qname
        case M.lookup dcon primMap of
          Just p  -> pure $ PrimAppE p []
          Nothing -> do
            -- Just a placeholder for now, the typechecker will fill this hole.
            ty <- newMetaTy
            pure $ DataConE ty dcon []
  
      -- TODO: timeit: parsing it's type isn't straightforward.
  
      InfixApp _ e1 (QVarOp _ (UnQual _ (Symbol _ "!!!"))) e2 -> do
        e1' <- desugarExp type_syns toplevel e1
        case e2 of
          Lit _ lit -> do
            let i = litToInt lit
            pure $ ProjE i e1'
          _ -> error $ "desugarExp: !!! expects a integer. Got: " ++ prettyPrint e2
  
      InfixApp _ e1 op e2 -> do
        e1' <- desugarExp type_syns toplevel e1
        e2' <- desugarExp type_syns toplevel e2
        case op of
          QVarOp _ (UnQual _ (Symbol _ "&")) -> do
            pure $ Ext (LinearExt (ReverseAppE e2' e1'))
          _ -> do
            let op' = desugarOp op
            pure $ PrimAppE op' [e1', e2']
  
      NegApp _ e1 -> do
        e1' <- desugarExp type_syns toplevel e1
        pure $ PrimAppE SubP [LitE 0, e1']
  
      _ -> error ("desugarExp: Unsupported expression: " ++ prettyPrint e)
  
  desugarFun :: (Show a,  Pretty a) => TypeSynEnv -> TopTyEnv -> TopTyEnv -> Decl a -> PassM (Var, [Var], TyScheme, Exp0)
  desugarFun type_syns toplevel env decl =
    case decl of
      FunBind _ [Match _ fname pats (UnGuardedRhs _ bod) _where] -> do
        let fname_str = nameToStr fname
            fname_var = toVar (fname_str)
        (vars, arg_tys,bindss) <- unzip3 <$> mapM (desugarPatWithTy type_syns) pats
        let binds = concat bindss
            args = vars
        fun_ty <- case M.lookup fname_var env of
                    Nothing -> do
                       ret_ty  <- newMetaTy
                       let funty = ArrowTy arg_tys ret_ty
                       pure $ (ForAll [] funty)
                    Just ty -> pure ty
        bod' <- desugarExp type_syns toplevel bod
        pure $ (fname_var, args, unCurryTopTy fun_ty, (mkLets binds bod'))
      _ -> error $ "desugarFun: Found a function with multiple RHS, " ++ prettyPrint decl
  
  multiArgsToOne :: [Var] -> [Ty0] -> Exp0 -> (Var, Exp0)
  multiArgsToOne args tys ex =
    let new_arg = toVar "multi_arg"
    in (new_arg, tuplizeRefs new_arg args tys ex)
  
  collectTopTy :: (Show a,  Pretty a) => TypeSynEnv -> Decl a -> TopTyEnv -> TopTyEnv
  collectTopTy type_syns d env =
    case d of
      TypeSig _ names ty ->
        let ty' = desugarTopType type_syns ty
        in foldr (\name acc ->
                    let tycon_var = toVar (nameToStr name) in
                    case M.lookup tycon_var acc of
                      Nothing ->  M.insert tycon_var ty' acc
                      Just{} -> error $ "collectTopTy: Multiple type signatures for: " ++ show tycon_var)
           env names
      _ -> env
  
  collectTypeSynonyms :: (Show a,  Pretty a) => TypeSynEnv -> Decl a -> TypeSynEnv
  collectTypeSynonyms env d =
    case d of
      TypeDecl _ (DHead _ name) ty ->
        let ty' = desugarType env ty
            tycon = nameToStr name
        in case M.lookup tycon env of
             Nothing -> M.insert tycon ty' env
             Just{} -> error $ "collectTypeSynonyms: Multiple type synonym declarations: " ++ show tycon
      _ -> env
  
  collectTopLevel :: (Show a,  Pretty a) => TypeSynEnv -> TopTyEnv -> Decl a -> PassM (Maybe TopLevel)
  collectTopLevel type_syns env decl =
    let toplevel = env in
    case decl of
      -- 'collectTopTy' takes care of this.
      TypeSig{} -> pure Nothing
  
      -- 'collectTypeSynonyms'.
      TypeDecl{} -> pure Nothing
  
      AnnPragma _ annotation -> 
        case annotation of 
              TypeAnn _ (Ident _ tycon) (Lit _ (String _ "Factored" _)) -> pure $ Just (MemLayoutTy tycon FullyFactored)
              TypeAnn _ (Ident _ tycon) (Lit _ (String _ "Linear" _)) -> pure $ Just (MemLayoutTy tycon Linear)
              _ -> error "Memory Layout not yet supported!"
  
  
      DataDecl _ (DataType _) _ctx decl_head cons _deriving_binds -> do
        let (ty_name,  ty_args) = desugarDeclHead decl_head
            cons' = map (desugarConstr type_syns) cons
        if ty_name `S.member` builtinTys
        then error $ sdoc ty_name ++ " is a built-in type."
        -- Default to Linear memory layout but we update it using Ann pragmas if available.
        else pure $ Just $ HDDef (DDef ty_name ty_args cons' Linear)
  
      -- Reserved for HS.
      PatBind _ (PVar _ (Ident _ "main")) (UnGuardedRhs _ _) _binds ->
        pure Nothing
  
      PatBind _ (PVar _ (Ident _ "gibbon_main")) (UnGuardedRhs _ rhs) _binds -> do
        rhs' <- fixupSpawn <$> verifyBenchEAssumptions True <$> desugarExp type_syns toplevel rhs
        ty <- newMetaTy
        pure $ Just $ HMain $ Just (rhs', ty)
  
      PatBind _ (PVar _ (Ident _ fn)) (UnGuardedRhs _ rhs) _binds ->
         case M.lookup (toVar fn) env of
           Nothing -> error $ "collectTopLevel: Top-level binding with no type signature: " ++ fn
           Just fun_ty ->
               -- This is a top-level function binding of the form:
               --     f = \x -> ...
               case rhs of
                 Lambda _ pats bod -> do
                   bod' <- desugarExp type_syns toplevel bod
                   case pats of
                     [] -> error "Impossible"
                     _  -> do
                       (vars,_tys,bindss) <- unzip3 <$> mapM (desugarPatWithTy type_syns) pats
                       let binds = concat bindss
                           args = vars
                       pure $ Just $ HFunDef (FunDef { funName = toVar fn
                                                     , funArgs = args
                                                     , funTy   = fun_ty
                                                     , funBody = fixupSpawn (mkLets binds bod')
                                                     , funMeta = FunMeta { funRec = NotRec
                                                                         , funInline = NoInline
                                                                         , funCanTriggerGC = False
                                                                         }
                                                     })
  
                 -- This is a top-level function that doesn't take any arguments.
                 _ -> do
                   rhs' <- desugarExp type_syns toplevel rhs
                   let fun_ty'  = ArrowTy [] (tyFromScheme fun_ty)
                       fun_ty'' = ForAll (tyVarsInTy fun_ty') fun_ty'
                   pure $ Just $ HFunDef (FunDef { funName = toVar fn
                                                 , funArgs = []
                                                 , funTy   = fun_ty''
                                                 , funBody = fixupSpawn rhs'
                                                 , funMeta = FunMeta { funRec = NotRec
                                                                     , funInline = NoInline
                                                                     , funCanTriggerGC = False
                                                                     }
                                                 })
  
  
      FunBind{} -> do (name,args,ty,bod) <- desugarFun type_syns toplevel env decl
                      pure $ Just $ HFunDef (FunDef { funName = name
                                                    , funArgs = args
                                                    , funTy   = ty
                                                    , funBody = fixupSpawn bod
                                                    , funMeta = FunMeta { funRec = NotRec
                                                                        , funInline = NoInline
                                                                        , funCanTriggerGC = False
                                                                        }
                                                    })
  
      InlineSig _ _ _ qname -> pure $ Just $ HInline (toVar $ qnameToStr qname)
  
      _ -> error $ "collectTopLevel: Unsupported top-level expression: " ++ show decl
  
  
  -- pure $ LitE (litToInt lit)
  desugarLiteral :: Literal a -> PassM Exp0
  desugarLiteral lit =
    case lit of
      (Int _ i _)  -> pure $ LitE (fromIntegral i)
      (Char _ chr _) -> pure $ CharE chr
      (Frac _ i _) -> pure $ FloatE (fromRational i)
      (String _ str _) -> do
        vec <- gensym (toVar "vec")
        let n = length str
            init_vec = LetE (vec,[],VectorTy CharTy, PrimAppE (VAllocP CharTy) [LitE n])
            fn i c b = LetE ("_",[],VectorTy CharTy,
                             PrimAppE (InplaceVUpdateP CharTy) [VarE vec, LitE i, CharE c])
                       b
            add_chars = foldr (\(i,chr) acc -> fn i chr acc) (VarE vec)
                          (reverse $ zip [0..n-1] str)
        pure $ init_vec add_chars
  
      _ -> error ("desugarLiteral: Only integer litrals are allowed: " ++ prettyPrint lit)
  
  
  litToInt :: Literal a -> Int
  litToInt (Int _ i _) = (fromIntegral i)
  litToInt lit         = error ("litToInt: Not an integer: " ++ prettyPrint lit)
  
  litToString :: Literal a -> String
  litToString (String _ a _) = a
  litToString lit            = error ("litToString: Expected a String, got: " ++ prettyPrint lit)
  
  qnameToStr :: H.QName a -> String
  qnameToStr qname =
    case qname of
      Qual _ mname n -> (mnameToStr mname ++ "." ++ nameToStr n)
      UnQual _ n     -> (nameToStr n)
      Special{}      -> error $ "qnameToStr: Special identifiers not supported: " ++ prettyPrint qname
  
  mnameToStr :: ModuleName a -> String
  mnameToStr (ModuleName _ s) = s
  
  desugarOp :: QOp a -> (Prim Ty0)
  desugarOp qop =
    case qop of
      QVarOp _ (UnQual _ (Symbol _ op)) ->
        case M.lookup op primMap of
          Just pr -> pr
          Nothing -> error $ "desugarOp: Unsupported binary op: " ++ show op
      op -> error $ "desugarOp: Unsupported op: " ++ prettyPrint op
  
  desugarAlt :: (Show a,  Pretty a) => TypeSynEnv -> TopTyEnv -> Alt a -> PassM (DataCon, [(Var,Ty0)], Exp0)
  desugarAlt type_syns toplevel alt =
    case alt of
      Alt _ (PApp _ qname ps) (UnGuardedRhs _ rhs) Nothing -> do
        let conName = qnameToStr qname
        desugarCase ps conName rhs
      Alt _ (PWildCard _) (UnGuardedRhs _ rhs) _b ->
        desugarCase [] "_default" rhs
      Alt _ _ GuardedRhss{} _ -> error "desugarExp: Guarded RHS not supported in case."
      Alt _ _ _ Just{}        -> error "desugarExp: Where clauses not allowed in case."
      Alt _ pat _ _           -> error $ "desugarExp: Unsupported pattern in case: " ++ prettyPrint pat
    where
      desugarCase ps conName rhs = do
        ps' <- mapM (\x -> case x of
                              PVar _ v -> (pure . toVar . nameToStr) v
                              PWildCard _ -> gensym "wildcard_"
                              _        -> error $ "desugarExp: Non-variable pattern in case." ++ show x)
                    ps
        rhs' <- desugarExp type_syns toplevel rhs
        ps'' <- mapM (\v -> (v,) <$> newMetaTy) ps'
        pure (conName, ps'', rhs')
  
  generateBind :: (Show a,  Pretty a) => TypeSynEnv -> TopTyEnv -> TopTyEnv -> Decl a -> Exp0 -> PassM (Exp0)
  generateBind type_syns toplevel env decl exp2 =
    case decl of
      -- 'collectTopTy' takes care of this.
      TypeSig{} -> pure exp2
      -- 'collectTypeSynonyms' takes care of this.
      TypeDecl{} -> pure exp2
      PatBind _ _ _ Just{}        -> error "generateBind: where clauses not allowed"
      PatBind _ _ GuardedRhss{} _ -> error "generateBind: Guarded right hand side not supported."
      PatBind _ (PTuple _ Boxed pats) (UnGuardedRhs _ rhs) Nothing -> do
        rhs' <- desugarExp type_syns toplevel rhs
        w <- gensym "tup"
        ty' <- newMetaTy
        let tupexp e = LetE (w,[],ty',rhs') e
            binds = reverse $ zip pats [0..]
        prjexp <- generateTupleProjs toplevel env binds (VarE w) exp2
        pure $ tupexp prjexp
      PatBind _ pat (UnGuardedRhs _ rhs) Nothing -> do
        rhs' <- desugarExp type_syns toplevel rhs
        w <- case pat of
               PVar _ v    -> pure $ toVar (nameToStr v)
               PWildCard _ -> gensym "wildcard_"
               _           -> error $ "generateBind: " ++ show pat
        ty' <- case M.lookup w env of
                 Nothing -> newMetaTy
                 Just (ForAll _ ty) -> pure ty
        pure $ LetE (w, [], ty', rhs') exp2
      FunBind{} -> do (name,args,ty,bod) <- desugarFun type_syns toplevel env decl
                      pure $ LetE (name,[], tyFromScheme ty, Ext $ LambdaE (zip args (inTys ty)) bod) exp2
      oth -> error ("generateBind: Unsupported pattern: " ++ prettyPrint oth)
  
  generateTupleProjs :: (Show a, Pretty a) => TopTyEnv -> TopTyEnv -> [(Pat a,Int)] -> Exp0 -> Exp0 -> PassM (Exp0)
  generateTupleProjs _toplevel _env [] _tup exp2 = pure exp2
  generateTupleProjs toplevel env ((p,n):pats) tup exp2 =
      case p of
          (PVar _ v) -> do
              let w = toVar (nameToStr v)
              go w
          -- Don't bind wildcards from patterns.
          (PWildCard _) -> do
            generateTupleProjs toplevel env pats tup exp2
  
          _ -> error $ "generateTupleProjs: Pattern not handled: " ++ prettyPrint p
  
    where
      go w = do
          ty' <- case M.lookup w env of
                     Nothing -> newMetaTy
                     Just (ForAll _ ty) -> pure ty
          let prjexp = LetE (w,[],ty',ProjE n tup) exp2
          generateTupleProjs toplevel env pats tup prjexp
  
  desugarConstr :: (Show a,  Pretty a) => TypeSynEnv -> QualConDecl a -> (DataCon,[(IsBoxed, Ty0)])
  desugarConstr type_syns qdecl =
    case qdecl of
      QualConDecl _ _tyvars _ctx (ConDecl _ name arg_tys) ->
        -- N.B. This is a type scheme only to make the types work everywhere else
        -- in code. However, we shouldn't actually quantify over any additional
        -- type variables here. We only support Rank-1 types.
        ( nameToStr name , map (desugarType' type_syns) arg_tys )
      _ -> error ("desugarConstr: Unsupported data constructor: " ++ prettyPrint qdecl)
  
  desugarDeclHead :: DeclHead a -> (Var, [TyVar])
  desugarDeclHead = go []
    where
      go acc decl_head =
        case decl_head of
          DHead _ name -> (toVar (nameToStr name), acc)
          DHParen _ dh -> go acc dh
          DHApp _ dh tyvar ->
              let (v,acc') = go acc dh
              in (v, acc' ++ [desugarTyVarBind tyvar])
          _ -> error ("collectTopLevel: Unsupported data declaration: " ++ prettyPrint decl_head)
  
  desugarTyVarBind :: TyVarBind a -> TyVar
  desugarTyVarBind (UnkindedVar _ name) = UserTv (toVar (nameToStr name))
  desugarTyVarBind v@KindedVar{} = error $ "desugarTyVarBind: Vars with kinds not supported yet." ++ prettyPrint v
  
  desugarPatWithTy :: (Show a, Pretty a) => TypeSynEnv -> Pat a -> PassM (Var, Ty0, [L0.Binds Exp0])
  desugarPatWithTy type_syns pat =
    case pat of
      (PParen _ p)        -> desugarPatWithTy type_syns p
      (PatTypeSig _ p ty) -> do (v,_ty,binds) <- desugarPatWithTy type_syns p
                                pure (v, desugarType type_syns ty, binds)
      (PVar _ n)          -> do ty <- newMetaTy
                                pure (toVar (nameToStr n), ty, [])
      (PWildCard _)       -> do v <- gensym "wildcard_"
                                ty <- newMetaTy
                                pure (v,ty,[])
      (PTuple _ Boxed pats) -> do (vars,tys,bindss) <- unzip3 <$> mapM (desugarPatWithTy type_syns) pats
                                  tup <- gensym "tup"
                                  let binds0 = concat bindss
                                      binds1 = map (\(v,ty,i) -> (v,[],ty,ProjE i (VarE tup))) (zip3 vars tys [0..])
                                      tupty = ProdTy tys
                                      -- current bindings: binds1, recursive bindings: binds0
                                  pure (tup,tupty,binds1 ++ binds0)
  
      (PApp _ (UnQual _ (Ident _ "Ur")) [one]) -> desugarPatWithTy type_syns one
  
      _ -> error ("desugarPatWithTy: Unsupported pattern: " ++ show pat)
  
  nameToStr :: Name a -> String
  nameToStr (Ident _ s)  = s
  nameToStr (Symbol _ s) = s
  
  instance Pretty SrcSpanInfo where
  
  -- | SpawnE's are parsed in a strange way. If we see a 'spawn (f x1 x2)',
  -- we parse it as 'SpawnE HOLE [] [(f x1 x2)]'. This function patches it
  -- to 'SpawnE f [] [x1 x2]'.
  fixupSpawn :: Exp0 -> Exp0
  fixupSpawn ex =
    case ex of
      Ext (LambdaE vars bod) -> Ext (LambdaE vars (go bod))
      Ext (PolyAppE a b)     -> Ext (PolyAppE (go a) (go b))
      Ext (FunRefE{})        -> ex
      Ext (BenchE fn tyapps args b) -> Ext (BenchE fn tyapps (map go args) b)
      Ext (ParE0 ls) -> Ext (ParE0 (map go ls))
      Ext (PrintPacked ty arg) -> Ext (PrintPacked ty (go arg))
      Ext (CopyPacked ty arg) -> Ext (CopyPacked ty (go arg))
      Ext (TravPacked ty arg) -> Ext (TravPacked ty (go arg))
      Ext (L0.L p e)    -> Ext (L0.L p (go e))
      Ext (LinearExt ext) ->
        case ext of
          ReverseAppE fn arg -> Ext (LinearExt (ReverseAppE (go fn) (go arg)))
          LseqE a b   -> Ext (LinearExt (LseqE (go a) (go b)))
          AliasE a    -> Ext (LinearExt (AliasE (go a)))
          ToLinearE a -> Ext (LinearExt (ToLinearE (go a)))
      -- Straightforward recursion ...
      VarE{}     -> ex
      LitE{}     -> ex
      CharE{}    -> ex
      FloatE{}   -> ex
      LitSymE{}  -> ex
      AppE fn tyapps args -> AppE fn tyapps (map go args)
      PrimAppE pr args -> PrimAppE pr (map go args)
      DataConE dcon tyapps args -> DataConE dcon tyapps (map go args)
      ProjE i e  -> ProjE i $ go e
      IfE a b c  -> IfE (go a) (go b) (go c)
      MkProdE ls -> MkProdE $ map go ls
      -- Only allow BenchE in tail position
      LetE (v,locs,ty,rhs) bod -> LetE (v,locs,ty, go rhs) (go bod)
      CaseE scrt mp -> CaseE (go scrt) $ map (\(a,b,c) -> (a,b, go c)) mp
      TimeIt e ty b -> TimeIt (go e) ty b
      WithArenaE v e -> WithArenaE v (go e)
      SpawnE _ _ args ->
        case args of
            [(AppE fn tyapps ls)] -> SpawnE fn tyapps ls
            _ -> error $ "fixupSpawn: incorrect use of spawn: " ++ sdoc ex
      SyncE   -> SyncE
      MapE{}  -> error $ "fixupSpawn: TODO MapE"
      FoldE{} -> error $ "fixupSpawn: TODO FoldE"
    where go = fixupSpawn
  
  -- | Verify some assumptions about BenchE.
  verifyBenchEAssumptions :: Bool -> Exp0 -> Exp0
  verifyBenchEAssumptions bench_allowed ex =
    case ex of
      Ext (LambdaE vars bod) -> Ext (LambdaE vars (not_allowed bod))
      Ext (PolyAppE a b)     -> Ext (PolyAppE (not_allowed a) (not_allowed b))
      Ext (FunRefE{})        -> ex
      Ext (BenchE _ tyapps args b) ->
        if bench_allowed then
          case args of
            ((VarE fn) : oth) -> Ext (BenchE fn tyapps oth b)
            _ -> error $ "desugarModule: bench is a reserved keyword. Usage: bench fn_name args. Got: " ++ sdoc args
        else error $ "verifyBenchEAssumptions: 'bench' can only be used as a tail of the main expression, but it was used in a function. In: " ++ sdoc ex
      Ext (ParE0 ls) -> Ext (ParE0 (map not_allowed ls))
      Ext (PrintPacked ty arg) -> Ext (PrintPacked ty (not_allowed arg))
      Ext (CopyPacked ty arg) -> Ext (CopyPacked ty (not_allowed arg))
      Ext (TravPacked ty arg) -> Ext (TravPacked ty (not_allowed arg))
      Ext (L0.L p e)    -> Ext (L0.L p (go e))
      Ext (LinearExt{}) -> error "verifyBenchEAssumptions: LinearExt not handled."
      -- Straightforward recursion ...
      VarE{}     -> ex
      LitE{}     -> ex
      CharE{}    -> ex
      FloatE{}   -> ex
      LitSymE{}  -> ex
      AppE fn tyapps args -> AppE fn tyapps (map not_allowed args)
      PrimAppE pr args -> PrimAppE pr (map not_allowed args)
      DataConE dcon tyapps args -> DataConE dcon tyapps (map not_allowed args)
      ProjE i e  -> ProjE i $ not_allowed e
      IfE a b c  -> IfE (not_allowed a) (go b) (go c)
      MkProdE ls -> MkProdE $ map not_allowed ls
      LetE (v,locs,ty,rhs) bod -> LetE (v,locs,ty, not_allowed rhs) (go bod)
      CaseE scrt mp -> CaseE (go scrt) $ map (\(a,b,c) -> (a,b, go c)) mp
      TimeIt e ty b -> TimeIt (not_allowed e) ty b
      WithArenaE v e -> WithArenaE v (go e)
      SpawnE fn tyapps args -> SpawnE fn tyapps (map not_allowed args)
      SyncE    -> SyncE
      MapE{}  -> error $ "verifyBenchEAssumptions: TODO MapE"
      FoldE{} -> error $ "verifyBenchEAssumptions: TODO FoldE"
    where go = verifyBenchEAssumptions bench_allowed
          not_allowed = verifyBenchEAssumptions False
  
  --------------------------------------------------------------------------------
  
  desugarLinearExts :: Prog0 -> PassM Prog0
  desugarLinearExts (Prog ddefs fundefs main) = do
      main' <- case main of
                 Nothing -> pure Nothing
                 Just (e,ty) -> do
                   let ty' = goty ty
                   e' <- go e
                   pure $ Just (e', ty')
      fundefs' <- mapM (\fn -> do
                             bod <- go (funBody fn)
                             let (ForAll tyvars ty) = (funTy fn)
                                 ty' = goty ty
                             pure $ fn { funBody = bod
                                       , funTy = (ForAll tyvars ty')
                                       })
                       fundefs
      pure (Prog ddefs fundefs' main')
    where
      goty :: Ty0 -> Ty0
      goty ty =
        case ty of
          ProdTy tys -> ProdTy (map goty tys)
          SymDictTy v t -> SymDictTy v (goty t)
          PDictTy k v -> PDictTy (goty k) (goty v)
          ArrowTy tys b  -> ArrowTy (map goty tys) (goty b)
          PackedTy "Ur" [one] -> one
          PackedTy t tys -> PackedTy t (map goty tys)
          VectorTy t -> VectorTy (goty t)
          ListTy t -> ListTy (goty t)
          _ -> ty
  
      go :: PreExp E0Ext Ty0 Ty0 -> PassM Exp0
      go ex =
        case ex of
          VarE{}    -> pure ex
          LitE{}    -> pure ex
          CharE{}   -> pure ex
          FloatE{}  -> pure ex
          LitSymE{} -> pure ex
          AppE f tyapps args -> do args' <- mapM go args
                                   pure (AppE f tyapps args')
          PrimAppE pr args   -> do args' <- mapM go args
                                   pure (PrimAppE pr args')
          LetE (v,locs,ty,rhs) bod -> do
            let ty' = goty ty
            rhs' <- go rhs
            bod' <- go bod
            pure $ LetE (v,locs,ty',rhs') bod'
          IfE a b c -> do a' <- go a
                          b' <- go b
                          c' <- go c
                          pure (IfE a' b' c')
          MkProdE ls -> do ls' <- mapM go ls
                           pure (MkProdE ls')
          ProjE i e  -> do e' <- go e
                           pure (ProjE i e')
          CaseE scrt alts -> do scrt' <- go scrt
                                alts' <- mapM (\(a,b,c) -> do c' <- go c
                                                              pure (a,b,c'))
                                              alts
                                pure (CaseE scrt' alts')
          DataConE _ "Ur" [arg]   -> do arg' <- go arg
                                        pure arg'
          DataConE locs dcon args -> do args' <- mapM go args
                                        pure (DataConE locs dcon args')
          TimeIt e ty b -> do e' <- go e
                              let ty' = goty ty
                              pure (TimeIt e' ty' b)
          WithArenaE v e -> do e' <- go e
                               pure (WithArenaE v e')
          SpawnE f tyapps args -> do args' <- mapM go args
                                     pure (SpawnE f tyapps args')
          SyncE -> pure SyncE
          MapE{}  -> error "desugarLinearExts: MapE"
          FoldE{} -> error "desugarLinearExts: FoldE"
          Ext ext ->
            case ext of
              LambdaE args bod -> do bod' <- go bod
                                     let args' = map (\(v,ty) -> (v,goty ty)) args
                                     pure (Ext (LambdaE args' bod'))
              PolyAppE fn arg  -> do fn' <- go fn
                                     arg' <- go arg
                                     pure (Ext (PolyAppE fn' arg'))
              FunRefE{} -> pure ex
              BenchE fn tyapps args b -> do args' <- mapM go args
                                            pure (Ext (BenchE fn tyapps args' b))
              ParE0 ls -> do ls' <- mapM go ls
                             pure (Ext (ParE0 ls'))
              PrintPacked ty arg -> do arg' <- go arg
                                       pure (Ext (PrintPacked ty arg'))
              CopyPacked ty arg -> do arg' <- go arg
                                      pure (Ext (CopyPacked ty arg'))
              TravPacked ty arg -> do arg' <- go arg
                                      pure (Ext (TravPacked ty arg'))
              L0.L p e -> do e' <- go e
                             pure (Ext (L0.L p e'))
              LinearExt lin ->
                case lin of
                  ReverseAppE fn (Ext (LinearExt (AliasE e))) -> do
                    fn' <- go fn
                    case fn' of
                      Ext (LambdaE [(v,ProdTy tys)] bod) -> do
                        let ty = Sf.headErr tys
                            bod'' = foldl' (\acc i -> gSubstE (ProjE i (VarE v)) (VarE v) acc) bod [0..(length tys)]
                        pure (LetE (v,[],ty,e) bod'')
                      _ -> error $ "desugarLinearExts: couldn't desugar " ++ sdoc ex
                  ReverseAppE fn arg -> do
                    fn'  <- go fn
                    arg' <- go arg
                    case fn' of
                      Ext (LambdaE [(v,ty)] bod) -> do
                        pure (LetE (v,[],ty,arg') bod)
                      _ -> error $ "desugarLinearExts: couldn't desugar " ++ sdoc ex
                  LseqE _ b   -> do b' <- go b
                                    pure b'
                  AliasE a    -> do v <- gensym "aliased"
                                    ty <- newMetaTy
                                    pure (LetE (v,[],ty,MkProdE [a,a]) (VarE v))
                  ToLinearE a -> do a' <- go a
                                    pure a'
  