content: {-# LANGUAGE OverloadedStrings #-}
  module Ermine.Pretty.G
    ( prettyG
    , defaultCxt
    ) where
  
  import Control.Lens hiding (index)
  import Control.Monad.State
  import Data.Foldable as Foldable
  import Data.Functor
  import Data.Functor.Rep
  import Data.Maybe
  import Data.Map as Map hiding (splitAt)
  import Data.Monoid
  import Data.Vector (Vector)
  import qualified Data.Vector as V
  import Data.Word
  import Ermine.Pretty
  import Ermine.Pretty.Core
  import Ermine.Syntax.Sort
  import Ermine.Syntax.G
  
  defaultCxt :: Sort -> Ref -> Doc
  defaultCxt U (Lit n) = int (fromIntegral n)
  defaultCxt s (Lit n) = "bad lit?" <+> text (show s) <> "{" <> int (fromIntegral n) <> "}"
  defaultCxt N (Native _) = "native"
  defaultCxt s (Native _) = "bad native?" <+> text (show s)
  defaultCxt s (Global n) = "global?" <+> text (show s) <> "{" <> prettyId n <> "}"
  defaultCxt s (Stack n) = "stack?" <+> text (show s) <> "{" <> int (fromIntegral n) <> "}"
  defaultCxt s (Local n) = "local?" <+> text (show s) <> "{" <> int (fromIntegral n) <> "}"
  
  prettyG :: [Doc] -> (Sort -> Ref -> Doc) -> G -> Doc
  prettyG vs pr (Case e bs) =
        "case"
    <+> prettyG vs pr e
    <+> "of"
    <+> prettyContinuation vs pr bs
  prettyG vs pr (CaseLit r bs) =
        "case#"
    <+> pr U r
    <+> "of"
    <+> prettyContinuation vs pr bs
  prettyG _  pr (App _n f xs) = hsep $ prettyFunc pr f : prettySorted (imap (fmap Foldable.toList . fmap . pr) xs)
  prettyG vs pr (Let bs e)    = prettyLet vs pr False bs e
  prettyG vs pr (LetRec bs e) = prettyLet vs pr True bs e
  prettyG _ _   Slot          = text "slot"
  
  prettyFunc :: (Sort -> Ref -> Doc) -> Func -> Doc
  prettyFunc pr (Ref  r) = pr B r
  prettyFunc _  (Con  t) = "<" <> text (show t) <> ">"
  
  rummage :: [a] -> Word64 -> Either a Word64
  rummage []     n = Right n
  rummage (x:_)  0 = Left x
  rummage (_:xs) n = rummage xs (n-1)
  
  push :: Sorted [Doc] -> (Sort -> Ref -> Doc) -> Sort -> Ref -> Doc
  push xs f s = either id (f s) . _Stack (rummage (xs^.sort s))
  
  prettyContinuation :: [Doc] -> (Sort -> Ref -> Doc) -> Continuation -> Doc
  prettyContinuation vs pr (Continuation bs d) = block $ (f <$> Map.toList bs) ++ ldd
    where
      ldd = maybeToList $ (\c -> "_ ->" <+> prettyG vs pr c) <$> d
      f (t, (n, c)) = "<" <> text (show t) <> ">"
                  <+> hsep (prettySorted ws ++ ["->" <+> prettyG rest (push ws pr) c])
        where (ws, rest) = splitSorted n vs
  
  prettySorted :: Sorted [Doc] -> [Doc]
  prettySorted cvs = ifoldMap go cvs where
    go cc ws = text (show cc) <> semiBraces ws <$ guard (not (Prelude.null ws))
  
  splitSorted :: Sorted Word64 -> [a] -> (Sorted [a], [a])
  splitSorted = runState . traverse (state . splitAt . fromIntegral)
  
  prettyLet :: [Doc] -> (Sort -> Ref -> Doc) -> Bool -> Vector PreClosure -> G -> Doc
  prettyLet vs pr rec bs e =
        "let"
    <+> block (zipWith (\w bd -> w <+> "=" <+> bd) ws (Foldable.toList bds))
    <+> "in"
    <+> prettyG rest pr' e
   where
   bl = V.length bs
   (ws, rest) = splitAt bl vs
   pr' = push (Sorted ws [] []) pr
   bds = prettyPreClosure rest (if rec then pr' else pr) <$> bs
  
  prettyPreClosure :: [Doc] -> (Sort -> Ref -> Doc) -> PreClosure -> Doc
  prettyPreClosure vs pr (PreClosure cfvs (LambdaForm fa ba up bo))
    = hsep
    $ dfvs : dl : prettySorted cbvs ++ ["->" <+> prettyG rest pr' bo]
    where
      dl = "\\" <> if up then "u" else "n"
      dfvs = semiBraces $ join $ Foldable.toList $ imap (fmap Foldable.toList . fmap . pr) cfvs
      (cbvs, rest) = splitSorted ba vs
      pr' c (Local n)
        | n < index fa c = pr c $ index cfvs c V.! fromIntegral n
        | otherwise = error "PANIC: prettyPreClosure: Bad variable reference"
        where
      pr' c r = push cbvs pr c r
  