content: module Gibbon.L1.GenSML where
  
  import Gibbon.L1.Syntax
  import Gibbon.Common
  
  import Text.PrettyPrint hiding ((<>))
  import Data.Maybe
  import Control.Monad
  import Data.Map hiding (foldr, fold, null, empty)
  import qualified Data.Map as Map
  import qualified Data.Set as Set
  import qualified Safe as Sf
  
  import Data.Foldable hiding ( toList )
  import Data.Graph
  import Data.Tree ( flatten )
  
  
  ppExt :: E1Ext () Ty1 -> Doc
  ppExt ext0 = case ext0 of
    BenchE _var _uts _pes _b -> error "BenchE"
    AddFixed _var _n -> error "AddFixed"
    StartOfPkdCursor _var -> error "StartOfPkdCursor"
  
  ppE :: Exp1 -> Doc
  ppE e0 = case e0 of
    VarE var -> ppVar var
    LitE n -> int n
    CharE c -> char c
    FloatE x -> double x
    LitSymE var -> doubleQuotes $ ppVar var
    AppE var _ pes -> ppAp (ppVar var) pes
    PrimAppE pr pes -> ppPrim pr pes
    LetE (v, _, _, e) pe' ->
      hsep
        [ "\n  let val", ppVar v, "="
        , ppE e, "in"
        , ppE pe', "end"
        ]
    IfE pe' pe2 pe3 ->
      ("\n  " <>) $ parens $ hsep
        [ "if", ppE pe'
        , "then", ppE pe2
        , "\n   else", ppE pe3
        ]
    MkProdE pes -> parens $ interleave comma $ ppE <$> pes
    ProjE 0 pe' -> parens $ hsep
      [ "case", ppE pe', "of"
      , "(t0, _) => t0"
      ]
    ProjE 1 pe' -> parens $ hsep
      [ "case", ppE pe', "of"
      , "(_, t1) => t1"
      ]
    ProjE n pe' -> parens $ hsep [hcat ["#", int $ succ n], ppE pe']
    CaseE pe' x0 ->
      parens $ hsep
        [ hsep ["case", ppE pe', "of"]
        , interleave "\n  |" ((\(dc, vs, e) -> hsep
          [ text dc
          , case vs of
            [] -> mempty
            _ -> parens $ interleave comma $ ppVar . fst <$> vs
          , "=>", ppE e
          ]) <$> x0)
        ]
    DataConE _ty0 s [] -> text s
    DataConE _ty0 s pes ->
      parens $ hsep
        [ text s
        , parens $ interleave comma $ ppE <$> pes
        ]
  
    TimeIt _pe' _ty0 _b -> _
    WithArenaE _var _pe' -> error "WithArenaE"
    SpawnE _var _ty0s _pes -> error "SpawnE"
    SyncE -> error "SyncE"
    MapE _x0 _pe' -> error "MapE"
    FoldE _x0 _x1 _pe' -> error "FoldE"
  
    Ext ee -> ppExt ee
  
  ppCurried :: Doc -> [Exp1] -> Doc
  ppCurried var pes = parens $ hsep $ var : (ppE <$> pes)
  
  ppAp :: Doc -> [Exp1] -> Doc
  ppAp var pes =
    parens $ var <> case pes of
      [] -> empty               -- don't confuse with application to unit (1 arg)
      [x] -> space <> ppE x
      _ -> parens $ interleave "," $ ppE <$> pes
  
  ppVar :: Var -> Doc
  ppVar = text . getVar
  
  getVar :: Var -> String
  getVar (Var s) = case unintern s of
    "val" -> "val_"
    "as" -> "as_"
    "open" -> "open_"
    "rec" -> "rec_"
    "fun" -> "fun_"
    "end" -> "end_"
    '_' : z -> "internal_" ++ z
    z -> z
  
  interleave :: Doc -> [Doc] -> Doc
  interleave sepr lst = case lst of
    [] -> mempty
    d : ds -> (d <+>) $ fold $ (sepr <+>) <$> ds
  
  binary :: String -> [Exp1] -> Doc
  binary opSym pes =
    parens $ hsep [l, text opSym, r]
    where
      (l, r) = extractBinary opSym pes
  
  extractBinary :: String -> [Exp1] -> (Doc, Doc)
  extractBinary opSym pes = case ppE <$> pes of
    [l, r] -> (l, r)
    es -> error $ fold
      [ "L0 error: (", opSym, ") is provided "
      , show $ length es, " arguments"
      ]
  
  extractUnary :: String -> [Exp1] -> Doc
  extractUnary opSym pes = case ppE <$> pes of
    [x] -> x
    es -> error $ fold
      [ "L0 error: (", opSym, ") is provided "
      , show $ length es, " arguments"
      ]
  
  ppFail :: String -> Doc
  ppFail s = hsep
    [ "raise"
    , parens $ hsep ["Fail", doubleQuotes $ text s]
    ]
  
  ppPrim :: Prim Ty1 -> [Exp1] -> Doc
  ppPrim pr pes = case pr of
    AddP -> binary "+" pes
    SubP -> binary "-" pes
    MulP -> binary "*" pes
    DivP -> binary "div" pes
    ModP -> binary "mod" pes
    ExpP -> binary "**" pes
    RandP -> ppCurried "MltonRandom.rand()" pes
    EqIntP -> binary "=" pes
    LtP -> binary "<" pes
    GtP -> binary ">" pes
    LtEqP -> binary "<=" pes
    GtEqP -> binary ">=" pes
    FAddP -> binary "+" pes
    FSubP -> binary "-" pes
    FMulP -> binary "*" pes
    FDivP -> binary "/" pes
    FExpP ->
      parens $ hsep
        [ "Math.pow"
        , parens $ hcat [l, comma, r]
        ]
      where
        (l, r) = extractBinary "pow" pes
    FRandP -> ppCurried "Random.randFloat" pes
    EqFloatP -> binary "=" pes
    EqCharP -> binary "=" pes
    FLtP -> binary "<" pes
    FGtP -> binary ">" pes
    FLtEqP -> binary "<=" pes
    FGtEqP -> binary ">=" pes
    FSqrtP -> ppAp "Math.sqrt" pes
    IntToFloatP -> ppAp "Real.fromInt" pes
    FloatToIntP -> ppAp "Int.fromReal" pes
    FTanP -> ppAp "Math.tan" pes
    EqSymP -> binary "=" pes
    EqBenchProgP _ -> error "GenSML: EqBenchProgP"
    OrP -> binary "orelse" pes
    AndP -> binary "andalso" pes
    MkTrue -> "true"
    MkFalse -> "false"
    ErrorP s _ -> ppFail s
    SizeParam -> int 1  -- ?
    IsBig -> error "IsBig"
    GetNumProcessors -> error "GetNumProcessors"
    PrintInt -> printer "Int" $ ppE $ Sf.headErr pes
    PrintChar -> printer "Char" $ ppE $ Sf.headErr pes
    PrintFloat -> printer "Float" $ ppE $ Sf.headErr pes
    PrintBool -> ppAp "(fn true => \"True\" | false => \"False\")" pes
    PrintSym -> ppAp "print" pes
    ReadInt -> error "ReadInt"  -- Have every program read from stdin?
    DictInsertP _ -> error "DictInsertP"
    DictLookupP _ -> error "DictLookupP"
    DictEmptyP _ -> error "DictEmptyP"
    DictHasKeyP _ -> error "DictHasKeyP"
    SymSetEmpty -> error "SymSetEmpty"
    SymSetInsert -> error "SymSetInsert"
    SymSetContains -> error "SymSetContains"
    SymHashEmpty -> error "SymHashEmpty"
    SymHashInsert -> error "SymHashInsert"
    SymHashLookup -> error "SymHashLookup"
    SymHashContains -> error "SymHashContains"
    IntHashEmpty -> error "IntHashEmpty"
    IntHashInsert -> error "IntHashInsert"
    IntHashLookup -> error "IntHashLookup"
    PDictAllocP _ty0 _ty0' -> error "PDictAllocP"
    PDictInsertP _ty0 _ty0' -> error "PDictInsertP"
    PDictLookupP _ty0 _ty0' -> error "PDictLookupP"
    PDictHasKeyP _ty0 _ty0' -> error "PDictHasKeyP"
    PDictForkP _ty0 _ty0' -> error "PDictForkP"
    PDictJoinP _ty0 _ty0' -> error "PDictJoinP"
    LLAllocP _ty0 -> error "LLAllocP"
    LLIsEmptyP _ty0 -> error "LLIsEmptyP"  -- Implement these? 
    LLConsP _ty0 -> error "LLConsP"
    LLHeadP _ty0 -> error "LLHeadP"
    LLTailP _ty0 -> error "LLTailP"
    LLFreeP _ty0 -> error "LLFreeP"
    LLFree2P _ty0 -> error "LLFree2P"
    LLCopyP _ty0 -> error "LLCopyP"
    VAllocP _ty0 ->
      ppAp "(fn internal__ => ArraySlice.full(Array.array(internal__, 0)))" pes
    VFreeP _ty0 -> error "VFreeP"
    VFree2P _ty0 -> error "VFree2P"
    VLengthP _ty0 -> ppAp "ArraySlice.length" pes
    VNthP _ty0 -> ppAp "ArraySlice.sub" pes
    VSliceP _ty0 -> case pes of
      [pe1, pe2, pe3] -> hcat
        [ "ArraySlice.subslice"
        , parens $ interleave comma
          [ ppE pe3
          , ppE pe1
          , parens $ "SOME" <+> ppE pe2
          ]
        ]
      _ -> _
    InplaceVUpdateP _ty0 -> hsep
        [ "let val _ ="
        , ppAp "ArraySlice.update" pes
        , "in", ppE $ Sf.headErr pes
        , "end"
        ]
    VConcatP _ty0 -> ppFail "VConcatP"
    VSortP _ty0 -> ppFail "VSortP"
    InplaceVSortP _ty0 -> ppCurried qsort pes
    VMergeP _ty0 -> ppFail "VMergeP"
    Write3dPpmFile _s -> error "Write3dPpmFile"
    ReadPackedFile _m_s _s _m_var _ty0 -> error "ReadPackedFile"
    WritePackedFile _s _ty0 -> error "WritePackedFile"
    ReadArrayFile _ma _ty0 -> error "ReadArrayFile"
    RequestEndOf -> error "RequestEndOf"
    RequestSizeOf -> error "RequestSizeOf"
    Gensym -> error "Gensym"
  
  ppProgram :: Prog1 -> Doc
  ppProgram prog = hcat
    [ ppDDefs $ ddefs prog
    , ppFunDefs $ fundefs prog
    , ppMainExpr $ mainExp prog
    , "\n"
    ]
  
  ppFunDefs :: Map Var (FunDef Var Exp1) -> Doc
  ppFunDefs funDefs =
    foldMap ppBlock organize
    where
      ppBlock = either ppValDef ppFunRec
      organize = sortDefs (elems funDefs) >>= separateDefs
  
  separateDefs :: [FunDef Var Exp1] -> [Either (FunDef Var Exp1) [FunDef Var Exp1]]
  separateDefs funDefs = case funDefs of
    [] -> []
    fd : fds -> case funArgs fd of
      [] -> Left fd : separateDefs fds
      _ -> case separateDefs fds of
        [] -> [Right [fd]]
        fds'@(Left _ : _) -> Right [fd] : fds'
        Right fds' : fds'' ->  Right (fd : fds') : fds''
  
  ppValDef :: FunDef Var Exp1 -> Doc
  ppValDef funDef =
    hsep
      [ "val"
      , ppVar $ funName funDef
      , "="
      , ppE $ funBody funDef
      ] <> semi
  
  ppFunRec :: [FunDef Var Exp1] -> Doc
  ppFunRec fdefs =
    reduceFunDefs "fun" (Sf.headErr fdefs) $
      foldr (reduceFunDefs "and") ";\n" (Sf.tailErr fdefs)
  
  reduceFunDefs :: Doc -> FunDef Var Exp1 -> Doc -> Doc
  reduceFunDefs keyword funDef doc =
    "\n" <> case funArgs funDef of
      [] -> hsep
        [ keyword
        , ppVar $ funName funDef
        , "="
        , ppE $ funBody funDef
        ] <> doc
      fargs -> hsep
        [ keyword
        , ppVar name
        , parens $ interleave comma $ ppVar <$> fargs
        , "="
        , case name of
          "print_check" -> parens mempty
          "print_space" -> "print \" \""
          "print_newline" -> "print \"\\n\""
          _ -> ppE $ funBody funDef
        ] <> doc
        where name = funName funDef
  
  ppMainExpr :: Maybe (Exp1, Ty1) -> Doc
  ppMainExpr opt = case opt of
    Nothing -> mempty
    Just (exp1, ty1) -> 
      "val _ = " <> printerTy1 ty1 (ppE exp1) <> semi
      <> "\nval _ = print \"\\n\""
  
  ppDDefs :: DDefs1 -> Doc
  ppDDefs ddefs = case elems ddefs of
    [] -> mempty
    h : t -> hsep
      [ "datatype"
      , ppDDef h
      , hcat $ ("\nand" <+>) . ppDDef <$> t
      , ";\n"
      ]
  
  ppDDef :: DDef1 -> Doc
  ppDDef ddef = hsep
    [ hsep $ ppTyVar <$> tyArgs ddef
    , ("dat_" <>) $ ppVar $ tyName ddef
    , "="
    , interleave
        "|"
        (ppBody <$> dataCons ddef)
    ]
    where
      ppBody (s, lst) = text s <+> case lst of
        [] -> mempty
        _ -> "of" <+> parens (interleave " *" $ ppTy1 . snd <$> lst)
  
  ppTyVar :: TyVar -> Doc
  ppTyVar tyVar = case tyVar of
    BoundTv var -> "'" <> ppVar var
    SkolemTv _s _n -> _
    UserTv var -> "'" <> ppVar var
  
  ppTy1 :: Ty1 -> Doc
  ppTy1 ty1 = case ty1 of
    IntTy -> "int"
    CharTy -> "char"
    FloatTy -> "real"
    BoolTy -> "bool"
    ProdTy ty1s -> interleave " * " $ ppTy1 <$> ty1s
    SymDictTy _m_var _ty1' -> _
    PDictTy _ty1' _ty12 -> _
    SymSetTy -> _
    SymHashTy -> _
    IntHashTy -> _
    PackedTy s () -> " dat_" <> text s
    VectorTy ty1' -> ppTy1 ty1' <+> "ArraySlice.slice"
    ListTy ty1' -> ppTy1 ty1' <+> "list"
    ArenaTy -> _
  
    SymTy -> _
    PtrTy -> _
    CursorTy -> _
    _ -> error "ppTy1: unexpected type"
  
  printerTy1 :: Ty1 -> Doc -> Doc
  printerTy1 ty1 d = case ty1 of
    IntTy -> printer "Int" d
    CharTy -> printer "Char" d
    FloatTy -> printer "Float" d
    SymTy -> _
    BoolTy -> parens $ "(fn true => print \"True\" | false => print \"False\") " <> d
    ProdTy [] -> "let val () = " <> d <> " in (print \"#()\") end"
    ProdTy uts -> 
      parens $ hsep
        [ "case", d, "of"
        , parens $ interleave comma $ ("x__" <>) . int . fst <$> zip [1..] uts
        , "-> let"
        , "val _ = print \"#(\""
        , foldMap ppSub $ zip [1..] uts
        , "val _ = print \")\""
        , "in ()"
        ]
      where
        ppSub (i, x) = hsep
          [ "val _ ="
          , printerTy1 x $ "x__" <> int i
          , "val _ = print \" \""
          ]
    SymDictTy _m_var _ut -> _
    PackedTy s () -> "internal_print_" <> text s <> parens d
    VectorTy ut -> 
      parens $ hsep
        [ quotePrint "#["
        , toss $ hsep 
          [ "case length", d, "of"
          , "0 -> ()"
          , "1 ->", printerTy1 ut $ "ArraySlice.sub" <> parens (d <> ", 0")
          , "_ ->"
          , toss $ printerTy1 ut $ "ArraySlice.sub" <> parens (d <> ", 0")
          , "ArraySlice.app", parens $ 
            "fn y__ => " <> printerTy1 ut "y__"
          , "xs__"
          ]
        , "print \"]\""
        ]
    PDictTy _ut _ut' -> _
    ListTy ut -> 
      parens $ hsep
        [ quotePrint "["
        , toss $ hsep [ "case", d, "of"
          , "[] -> ()"
          , "| [x__] ->", printerTy1 ut "x__"
          , "| [x__ :: xs__] ->"
          , toss $ printerTy1 ut "x__"
          , "list.app", parens $ 
            "fn y__ => " <> quotePrint ", " <> printerTy1 ut "y__"
          , "xs__"
          ]
        , "print \"]\""
        ]
    ArenaTy -> _
    SymSetTy -> _
    SymHashTy -> _
    IntHashTy -> _
    PtrTy -> _
    CursorTy -> _
    _ -> error "printerTy1: unexpected type"
  
  printer :: Doc -> Doc -> Doc
  printer p d = parens $ "print" <> parens (p <> ".toString" <> parens d)
  
  toss :: Doc -> Doc
  toss s = "let val _ = " <> s <> " in "
  justPrint :: Doc -> Doc
  justPrint s = toss $ "print " <> s
  quotePrint :: Doc -> Doc
  quotePrint = justPrint . quotes
  
  qsort :: Doc
  qsort = parens $ text
    "fn arr => fn cmp => \n\
    \  let\n\
    \    fun qsort(arr, lo, hi) = \n\
    \      if cmp lo hi < 0 then\n\
    \        let\n\
    \          val pivot = ArraySlice.sub(arr, hi)\n\
    \          val i = ref (lo - 1)\n\
    \          val j = ref lo\n\
    \          val _ = \n\
    \            while cmp (!j) (hi - 1) < 1 do\n\
    \              let\n\
    \                val _ = \n\
    \                  if cmp (ArraySlice.sub(arr, !j)) pivot < 0 then\n\
    \                    let\n\
    \                      val _ = i := !i + 1\n\
    \                      val tmp = ArraySlice.sub(arr, !i)\n\
    \                      val _ = ArraySlice.update(arr, !i, ArraySlice.sub(arr, !j))\n\
    \                      val _ = ArraySlice.update(arr, !j, tmp)\n\
    \                    in\n\
    \                      ()\n\
    \                    end\n\
    \                  else ()\n\
    \              in\n\
    \                j := !j + 1\n\
    \              end\n\
    \          val tmp = ArraySlice.sub(arr, !i + 1)\n\
    \          val _ = ArraySlice.update(arr, !i + 1, ArraySlice.sub(arr, hi))\n\
    \          val _ = ArraySlice.update(arr, hi, tmp)\n\
    \          val p = !i + 1\n\
    \          val _ = qsort(arr, lo, p - 1)\n\
    \          val _ = qsort(arr, p + 1, hi)\n\
    \        in\n\
    \          ()\n\
    \        end\n\
    \    else ()\n\
    \    val _ = qsort(arr, 0, ArraySlice.length arr - 1)\n\
    \  in\n\
    \    arr\
    \  end\n"
  
  varsEs :: Set.Set String -> [Exp1] -> Set.Set String
  varsEs = foldMap . varsE
  
  varsE :: Set.Set String -> Exp1 -> Set.Set String
  varsE vs pe0 = case pe0 of
    -- VarE var -> collect var
    VarE _ -> mempty
    AppE var _ pes -> vpes pes <> collect var
    PrimAppE _ pes -> vpes pes
    LetE (_, _, _, pe') pe -> vpe pe <> vpe pe'
    IfE pe pe' pe3 -> vpes [pe, pe', pe3]
    MkProdE pes -> vpes pes
    ProjE _ pe -> vpe pe
    CaseE pe x0 -> vpe pe <> foldMap (\(_, _, pe') -> vpe pe') x0
    DataConE _ _ pes -> vpes pes
    TimeIt pe _ _ -> vpe pe
    WithArenaE _ pe -> vpe pe
    SpawnE _ _ pes -> vpes pes
    SyncE -> _
    MapE _ _ -> _
    FoldE {} -> _
    _ -> mempty
    where
      vpe = varsE vs
      vpes = varsEs vs
      collect var
        | Set.member s vs = Set.singleton s
        | otherwise = mempty
        where s = getVar var
  
  addFunBinding :: FunDef Var ex -> Map String (FunDef Var ex) -> Map String (FunDef Var ex)
  addFunBinding funDef = Map.insert (getVar $ funName funDef) funDef
  
  allFunEntries :: [FunDef Var ex] -> Map String (FunDef Var ex)
  allFunEntries = foldr addFunBinding Map.empty
  
  allFunNames :: [FunDef Var ex] -> Set.Set String
  allFunNames = Set.fromList . fmap (getVar . funName)
  
  getDependencies :: [FunDef Var Exp1] -> Map String [FunDef Var Exp1]
  getDependencies funDefs =
    foldr reduceDeps Map.empty funDefs
    where
      funMap = allFunEntries funDefs
      funSet = allFunNames funDefs
      toNode = fromMaybe _ . flip Map.lookup funMap
      toDep = fmap toNode . Set.toList . varsE funSet . funBody
      reduceDeps = insert . getVar . funName <*> toDep
  
  sortDefs :: [FunDef Var Exp1] -> [[FunDef Var Exp1]]
  sortDefs defs =
    fmap ((\(_, n, _) -> n) . back) . flatten <$> scc gr
    where
      (gr, back, _) = graphFromEdges $ mkNode <$> toList depMap
      mkNode (s, lst) = (s, fromMaybe _ (Map.lookup s nameMap), lst)
      depMap = getDependencies defs
      nameMap = fromList $ join ((,) . getVar . funName) <$> defs
  