Language.hs

module Language where
import Utils

import Prelude hiding (seq)

data Expr a
   = EVar Name               -- ^ Variables
   | ENum Int                -- ^ Numbers
   | EConstr Int Int         -- ^ Constructor tag arity
   | EAp (Expr a) (Expr a)   -- ^ Applications
   | ELet                    -- ^ Let(rec) expressions
       IsRec                 -- ^   boolean with True = recursive,
       [(a, Expr a)]         -- ^   Definitions
       (Expr a)              -- ^   Boy of let(rec)
   | ECase                   -- ^ Case expression
       (Expr a)              -- ^   Expression to scrutinise
       [Alter a]             -- ^   Alternatives
   | ELam [a] (Expr a)       -- ^ Lambda abstraction
   deriving Show

type CoreExpr = Expr Name

type Name = String

type IsRec = Bool
recursive, nonRecursive :: IsRec
recursive    = True
nonRecursive = False

bindersOf :: [(a,b)] -> [a]
bindersOf defns =  [name | (name, _rhs) <- defns]

rhssOf       :: [(a, b)] -> [b]
rhssOf defns =  [rhs | (_name, rhs) <- defns]

type Alter a = (Int, [a], Expr a)
type CoreAlt = Alter Name

isAtomicExpr :: Expr a -> Bool
isAtomicExpr (EVar _v) = True
isAtomicExpr (ENum _n) = True
isAtomicExpr _e        = False

type Program a = [ScDefn a]
type CoreProgram = Program Name

type ScDefn a = (Name, [a], Expr a)
type CoreScDefn = ScDefn Name

sample1 :: CoreProgram
sample1 =
  [("main",   [],    (EAp (EVar "double") (ENum 21))),
   ("double", ["x"], (EAp (EAp (EVar "+") (EVar "x")) (EVar "x")))
  ]

prelude :: String
prelude =
  unlines
  [ "I x = x ;"
  , "K x y = x ;"
  , "K1 x y = y ;"
  , "S f g x = f x (g x) ;"
  , "compose f g x = f (g x) ;"
  , "twice f = compose f f"
  ]

preludeDefs :: CoreProgram
preludeDefs
 = [ ("I", ["x"], EVar "x")
   , ("K", ["x","y"], EVar "x")
   , ("K1",["x","y"], EVar "y")
   , ("S", ["f","g","x"], EAp (EAp (EVar "f") (EVar "x"))
                              (EAp (EVar "g") (EVar "x")))
   , ("compose", ["f","g","x"], EAp (EVar "f")
                                     (EAp (EVar "g") (EVar "x")))
   , ("twice", ["f"], EAp (EAp (EVar "compose") (EVar "f")) (EVar "f"))
   ]

-----

mkMultiAp :: Int -> CoreExpr -> CoreExpr -> CoreExpr
mkMultiAp n e1 e2 = foldl EAp e1 $ replicate n e2
-- mkMultiAp n e1 e2 = foldl EAp e1 (take n e2s)
--                    where
--                    e2s = e2 : e2s

class Iseq iseq where
  iNil :: iseq
  iStr :: String -> iseq
  iAppend :: iseq -> iseq -> iseq
  iNewline :: iseq
  iIndent :: iseq -> iseq
  iDisplay :: iseq -> String

infixr 5 `iAppend`

------

data Fixity
    = L | N | R
    deriving (Eq, Show)

ops :: [(Name, (Int, Fixity))]
ops = [ ("*", (5, R)), ("/", (5, N))
      , ("+", (4, R)), ("-", (4, N))
      , ("==", (3, N)), ("~=", (3, N))
      , (">", (3, N)), (">=", (3, N)), ("<", (3, N)), ("<=", (3, N))
      , ("&", (2, R))
      , ("|", (1, R)) ]

pprExpr :: (Int, Fixity) -> CoreExpr -> IseqRep
pprExpr _ (EVar v) = iStr v
pprExpr _ (ENum n) = iStr $ show n
pprExpr _ (EConstr tn a)
  = iConcat [iStr "Pack{", iStr (show tn), iStr ",", iStr (show a), iStr "}"]
pprExpr (cpr, cas) (EAp (EAp (EVar op) e1) e2)
  | Just (f@(p, a)) <- op `lookup` ops
  , let unparened =
          case a of
          L -> iConcat [pprExpr f e1, iStr " ", iStr op, iStr " ", pprExpr (p, N) e2]
          R -> iConcat [pprExpr (p, N) e1, iStr " ", iStr op, iStr " ", pprExpr f e2]
          N -> iConcat [pprExpr f e1, iStr " ", iStr op, iStr " ", pprExpr f e2]
        parened = iConcat [iStr "(", unparened, iStr ")"]
        result
          | p >  cpr                           =  unparened
          | p == cpr && cas == a && cas /= N   =  unparened
          | p == cpr && cas == a  {-cas == N-} =  parened
          | p == cpr  {-cas /= a-}             =  parened
          | {-p < cpr-} otherwise              =  parened
  = result
pprExpr _ (EAp e1 e2)
  = iConcat [pprExpr (6, L) e1, iStr " ", pprExpr (6, L) e2]
pprExpr _ (ELet isrec defns expr)
  = iIndent $
    iConcat [ iStr keyword, iNewline
            , iStr "  ",iIndent (pprDefns defns),iNewline
            , iStr "in ",pprExpr (0, N) expr]
    where
    keyword | not isrec = "let"
            | isrec = "letrec"
pprExpr _ (ECase e as)
  = iConcat [ iStr "case", iStr " ", pprExpr (0, N) e, iStr " of "
            , iStr "    ", iIndent $ iInterleave (iStr ";" `iAppend` iNewline) $ map pprAlter as
            ]
    where
    pprAlter (tn, ns, ae)
      = iConcat [ iInterleave (iStr " ") $ iStr ("<" ++ show tn ++ ">") : map iStr ns
                , iStr " -> ", pprExpr (0, N) ae
                ]
pprExpr _ (ELam ns e)
  = iConcat $ [iInterleave (iStr " ") $ map iStr $ "\\" : ns, iStr " ", pprExpr (0, N) e]

pprAExpr :: CoreExpr -> IseqRep
pprAExpr e
  | isAtomicExpr e = pprExpr (0, N) e
  | otherwise = iStr "(" `iAppend` pprExpr (0, N) e `iAppend` iStr ")"

pprDefns :: [(Name, CoreExpr)] -> IseqRep
pprDefns defns = iInterleave sep (map pprDefn defns)
                 where
                 sep = iConcat [ iStr ";", iNewline ]

pprDefn :: (Name, CoreExpr) -> IseqRep
pprDefn (name, expr)
  = iConcat [ iStr name, iStr " = ", iIndent (pprExpr (0, N) expr) ]

pprScDefn :: CoreScDefn -> IseqRep
pprScDefn (name, ns, e)
  = iInterleave (iStr " ") (map iStr $ name : ns) `iAppend`
    iStr " = " `iAppend` pprExpr (0, N) e `iAppend` iNewline

iInterleave :: Iseq a => a -> [a] -> a
iInterleave sep = irec
  where
    irec []     =  iNil
    irec [x]    =  x
    irec (x:xs) =  x `iAppend` sep `iAppend` irec xs

iConcat :: Iseq a => [a] -> a
iConcat = foldr iAppend iNil

pprint :: CoreProgram -> String
pprint prog = iDisplay (pprProgram prog)

pprProgram :: CoreProgram -> IseqRep
pprProgram = iInterleave iNewline . map pprScDefn

data IseqRep
   = INil
   | IStr String
   | IAppend IseqRep IseqRep
   | IIndent IseqRep
   | INewline
   deriving Show

instance Iseq IseqRep where
  iNil              = INil
  iAppend INil seq2 = seq2
  iAppend seq1 INil = seq1
  iAppend seq1 seq2 = IAppend seq1 seq2
  iStr ""           = INil
  iStr str          = case break (== '\n') str of
    (_, [])        -> IStr str
    (w, _:str1)    -> IStr w `iAppend` iNewline `iAppend` iStr str1
  iIndent seq       = IIndent seq
  iNewline          = INewline
  iDisplay seq      = flatten 0 [(seq,0)]

flatten1 :: [IseqRep] -> String
flatten1 [] = ""
flatten1 (INil : seqs) = flatten1 seqs
flatten1 (IStr s : seqs) = s ++ (flatten1 seqs)
flatten1 (IAppend seq1 seq2 : seqs) = flatten1 (seq1 : seq2 : seqs)
flatten1 (_ : _) = error "flatten1: not implemented"

flatten :: Int -> [(IseqRep, Int)] -> String
flatten _    [] = ""
flatten col  ((INil, _) : seqs) = flatten col seqs
flatten col  ((IStr s, _): seqs) = s ++ (flatten (col + length s) seqs)
flatten col  ((IAppend seq1 seq2, indent) : seqs)
  = flatten col ((seq1, indent) : (seq2, indent) : seqs)
flatten _col ((INewline, indent) : seqs)
  = '\n' : (space indent) ++ (flatten indent seqs)
flatten col  ((IIndent seq, _indent) : seqs)
  = flatten col ((seq, col) : seqs)

space :: Int -> String
space n = replicate n ' '

examples :: CoreProgram
examples =
  [ ("f", ["x"],
     ELet True
      [ ("y", EAp (EAp (EVar "+") (EVar "x")) (ENum 1))
      , ("z", EAp (EAp (EVar "+") (EVar "y")) (ENum 1))]
      (EVar "z")),
    ("g", ["x"],
      EAp (EAp (EVar "+") (EVar "x")) (ENum 1)),

    -- h = x + y > p * length xs
    ("h", [],
     EAp
     (EAp
      (EVar ">")
      (EAp (EAp (EVar "+") (EVar "x")) (EVar "y")))
     (EAp
      (EAp
       (EVar "*")
       (EVar "p"))
      (EAp (EVar "length") (EVar "xs")))),

    -- u = 5 / (3 * x * (7 - (2 + x + 1)))
    ("u", [],
     ENum 5 |/|
     (ENum 3 |*| EVar "x" |*| (ENum 7 |-| (ENum 2 |+| EVar "x" |+| ENum 1)))),

    ("v", [],
     (ENum 1 |+| ENum 2) |+| ENum 3)
  ]
  where
    ap2 f x y = EAp (EAp f x) y
    (|*|) = ap2 (EVar "*")
    (|/|) = ap2 (EVar "/")
    (|+|) = ap2 (EVar "+")
    (|-|) = ap2 (EVar "-")
    infixr 5 |*|
    infix  5 |/|
    infixr 4 |+|
    infix  4 |-|


---------- old codes -----------

pprExprS :: CoreExpr -> String
pprExprS (ENum n) = show n
pprExprS (EVar v) = v
pprExprS (EAp e1 e2) = pprExprS e1 ++ " " ++ pprAExprS e2

pprAExprS :: CoreExpr -> String
pprAExprS e
  | isAtomicExpr e = pprExprS e
  | otherwise = "(" ++ pprExprS e ++ ")"

-----

pprExprN :: CoreExpr -> IseqRep
pprExprN (EVar v) = iStr v
pprExprN (ENum n) = iStr $ show n
pprExprN (EConstr tn a)
  = iConcat [iStr "Pack{", iStr (show tn), iStr ",", iStr (show a), iStr "}"]
pprExprN (EAp (EAp (EVar op) e1) e2)
  | op `elem` ["*", "/", "+", "-"
              , "==", "~=", ">", ">=", "<", "<="]
  = iConcat [pprAExprN e1, iStr " ", iStr op, iStr " ", pprAExprN e2]
pprExprN (EAp e1 e2)
  = iConcat [pprExprN e1, iStr " ", pprAExprN e2]
pprExprN (ELet isrec defns expr)
  = iIndent $
    iConcat [ iStr keyword, iNewline
            , iStr "  ",iIndent (pprDefnsN defns),iNewline
          , iStr "in ",pprExprN expr]
    where
    keyword | not isrec = "let"
            | isrec = "letrec"
pprExprN (ECase e as)
  = iConcat [ iStr "case", iStr " ", pprExprN e, iStr " of "
            , iStr "    ", iIndent $ iInterleave (iStr ";" `iAppend` iNewline) $ map pprAlter as
            ]
    where
    pprAlter (tn, ns, ae)
      = iConcat [ iInterleave (iStr " ") $ iStr ("<" ++ show tn ++ ">") : map iStr ns
                , iStr " -> ", pprExprN ae
                ]
pprExprN (ELam ns e)
  = iConcat $ [iInterleave (iStr " ") $ map iStr $ "\\" : ns, iStr " ", pprExprN e]

pprAExprN :: CoreExpr -> IseqRep
pprAExprN e
  | isAtomicExpr e = pprExprN e
  | otherwise = iStr "(" `iAppend` pprExprN e `iAppend` iStr ")"

pprDefnsN :: [(Name, CoreExpr)] -> IseqRep
pprDefnsN defns = iInterleave sep (map pprDefnN defns)
                 where
                 sep = iConcat [ iStr ";", iNewline ]

pprDefnN :: (Name, CoreExpr) -> IseqRep
pprDefnN (name, expr)
  = iConcat [ iStr name, iStr " = ", iIndent (pprExprN expr) ]

pprScDefnN :: CoreScDefn -> IseqRep
pprScDefnN (name, ns, e)
  = iInterleave (iStr " ") (map iStr $ name : ns) `iAppend`
    iStr " = " `iAppend` pprExprN e `iAppend` iNewline