-- | A module for lambda expressions

module Feldspar.DSL.Lambda where



import Control.Monad.State
import Data.Tree
import Data.Typeable

import Feldspar.DSL.Expression



-- | Unique identifier
type Ident = String

-- | Extensible lambda expressions
data Lam expr role a
  where
    Variable :: Ident -> Lam expr role a

    Value :: a -> Lam expr role a

    Lambda
        :: (Typeable rb, Typeable b)
        => (Lam expr ra a -> Lam expr rb b) -> Lam expr (ra -> rb) (a -> b)

    (:$:)
        :: (Typeable ra, Typeable a)
        => Lam expr (ra -> rb) (a -> b) -> Lam expr ra a -> Lam expr rb b
      -- Application. Using an infix operator makes it *a lot* easier to work
      -- with the library.

    Let :: String -> Lam expr (ra -> (ra -> rb) -> rb) (a -> (a -> b) -> b)

    Inject :: expr role a -> Lam expr role a

      -- Note: Using an infix operator for application makes it *a lot* easier
      -- to work with the library.

      -- TODO 'Value' is used for evaluating 'Lambda' expressions. It should not
      --      be exported to the user. This is only a temporary solution.

      -- TODO 'Lambda' should have a base name field instead of 'Let'.



-- | Let binding
let_ :: (Typeable ra, Typeable a, Typeable rb, Typeable b)
    => String  -- ^ Preferred base name
    -> Lam expr ra a -> (Lam expr ra a -> Lam expr rb b) -> Lam expr rb b
let_ base a f = Let base :$: a :$: Lambda f

instance ExprEq expr => ExprEq (Lam expr)
  where
    exprEq a b = evalState (exprEqLam a b) 0

instance ExprEq expr => Eq (Lam expr role a)
  where
    (==) = exprEq

freshVar
    :: String  -- ^ Base name
    -> State Integer (Lam expr role a)
freshVar base = do
    v <- get
    put (v+1)
    return $ Variable (base ++ show v)

exprEqLam :: ExprEq expr => Lam expr ra a -> Lam expr rb b -> State Integer Bool
exprEqLam (Variable i1) (Variable i2) = return (i1 == i2)
exprEqLam (Lambda f1) (Lambda f2) = do
    i <- get
    v1 <- freshVar ""
    put i
    v2 <- freshVar ""
    exprEqLam (f1 v1) (f2 v2)
exprEqLam (f1 :$: a1) (f2 :$: a2) = do
    aCond <- exprEqLam f1 f2
    if aCond
      then exprEqLam a1 a2
      else return False
exprEqLam (Inject a) (Inject b) = return (exprEq a b)
exprEqLam (Let _)    (Let _)    = return True
exprEqLam _ _ = return False
  -- This case includes 'Value', which is only supposed to be used during
  -- evaluation.

instance Eval expr => Eval (Lam expr)
  where
    eval (Variable ident) = error $ "Evaluating variable " ++ show ident
    eval (Value a)        = a
    eval (Lambda f)       = eval . f . Value
    eval (f :$: a)        = eval f $ eval a
    eval (Inject a)       = eval a
    eval (Let _)          = flip ($)

instance ExprShow expr => ExprShow (Lam expr)
  where
    exprShow = flip evalState 0 . exprShowLam

instance ExprShow (Lam expr) => Show (Lam expr role a)
  where
    show = exprShow

-- | Shallow application. Function argument must be a 'Lambda'.
shallowApply :: Lam expr (ra -> rb) (a -> b) -> Lam expr ra a -> Lam expr rb b
shallowApply (Lambda f) = f

infixr 0 $$
($$) = shallowApply

isVar :: Lam expr role a -> Bool
isVar (Variable _) = True
isVar _            = False

isLet :: Lam expr role a -> Bool
isLet (Let _ :$: _ :$: _) = True
isLet _                   = False

-- | Parser for infix operators of the form @"(op)"@
viewInfix :: String -> Maybe String
viewInfix ('(':op)
    | (')':op') <- reverse op = Just op'
viewInfix _ = Nothing

-- | Shows a partially applied expression
exprShowApp :: ExprShow expr
    => [String]              -- ^ Missing arguments
    -> Lam expr role a       -- ^ Partially applied expression
    -> State Integer String
exprShowApp args (f :$: a) = do
    aStr <- exprShowLam a
    exprShowApp (aStr:args) f
exprShowApp args f = do
    fStr <- exprShowLam f
    return $ case (viewInfix fStr, args) of
        (Just op, [a,b]) -> "(" ++ unwords [a,op,b] ++ ")"
        _                -> "(" ++ unwords (fStr : args) ++ ")"

exprShowLam :: ExprShow expr => Lam expr role a -> State Integer String
exprShowLam (Variable ident) = return ident
exprShowLam (Value _)        = error "exprShowLam: illegal use of Value"
exprShowLam (Lambda f) = do
    v@(Variable ident) <- freshVar "v"
    body <- exprShowLam (f v)
    return $ "(\\" ++ ident ++ " -> "  ++ body ++ ")"
exprShowLam (Let base :$: a :$: Lambda f) = do
    v@(Variable ident) <- freshVar base
    aStr <- exprShowLam a
    body <- exprShowLam (f v)
    return $ "(let " ++ ident ++ " = " ++ aStr ++ " in "  ++ body ++ ")"
exprShowLam (f :$: a) = do
    aStr <- exprShowLam a
    exprShowApp [aStr] f
exprShowLam (Inject a) = return (exprShow a)

-- | Converts a partially applied expression to a tree
lamToTreeApp :: ExprShow expr
    => Forest String                -- ^ Missing arguments
    -> Lam expr role a              -- ^ Partially applied expression
    -> State Integer (Tree String)
lamToTreeApp args (f :$: a) = do
    aTree <- lamToTree a
    lamToTreeApp (aTree : args) f
lamToTreeApp args f = do
    fTree <- lamToTree f
    return $ Node "Apply" (fTree : args)

-- | Converts a lambda expression to a tree
lamToTree :: ExprShow expr => Lam expr role a -> State Integer (Tree String)
lamToTree (Variable ident) = return $ Node ("Variable " ++ ident) []
lamToTree (Value a)        = return $ Node "Value ..." []
lamToTree (Lambda f) = do
    v@(Variable ident) <- freshVar "v"
    body <- lamToTree (f v)
    return $ Node ("Lambda " ++ ident) [body]
lamToTree (Let base :$: a :$: Lambda f) = do
    v@(Variable ident) <- freshVar base
    aTree <- lamToTree a
    body  <- lamToTree (f v)
    return $ Node ("Let " ++ ident) [aTree,body]
lamToTree (f :$: a) = do
    aTree <- lamToTree a
    lamToTreeApp [aTree] f
lamToTree (Let base) = return $ Node ("Let " ++ base) []
lamToTree (Inject a) = return $ Node (exprShow a) []

-- | Show a lambda expression as a tree
showLamTree :: ExprShow expr => Lam expr role a -> String
showLamTree = drawTree . flip evalState 0 . lamToTree

-- | Print a lambda expression as a tree
drawLambda :: ExprShow expr => Lam expr role a -> IO ()
drawLambda = putStrLn . showLamTree