-- ----------------------------------------------------------------------
-- Copyright 2010-2011 National University of Ireland.
-- ----------------------------------------------------------------------
-- This file is part of DysVunctional Language.
-- 
-- DysVunctional Language is free software; you can redistribute it and/or modify
-- it under the terms of the GNU Affero General Public License as
-- published by the Free Software Foundation, either version 3 of the
--  License, or (at your option) any later version.
-- 
-- DysVunctional Language is distributed in the hope that it will be useful,
-- but WITHOUT ANY WARRANTY; without even the implied warranty of
-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-- GNU General Public License for more details.
-- 
-- You should have received a copy of the GNU Affero General Public License
-- along with DysVunctional Language.  If not, see <http://www.gnu.org/licenses/>.
-- ----------------------------------------------------------------------

{-# LANGUAGE NoImplicitPrelude, DeriveDataTypeable #-}
module FOL.Language.Expression where

import FOL.Language.Common
import FOL.Language.Pretty

import Data.Data
import Data.List
import Data.Monoid
import Data.Foldable (Foldable, foldMap)
import Data.Traversable

import Data.Generics.Uniplate.Data ()

import Control.Applicative

data Prog = Prog [Defn] Expr
            deriving (Eq, Show, Typeable, Data)

data Defn = Defn ShapedName [ShapedName] Expr
            deriving (Eq, Show, Typeable, Data)

data Shape
    = NilSh
    | RealSh
    | BoolSh
    | ConsSh Shape Shape
    | VectorSh [Shape]
    | ValuesSh [Shape]
    -- This is a hack.
    | FunctionSh
      deriving (Eq, Ord, Show, Typeable, Data)

data Expr
    = Var Name
    -- Literals
    | Nil
    | Bool Bool
    | Real Real
    -- Special forms
    | If Expr Expr Expr
    | Let (Bindings Name Expr) Expr
    | LetValues (Bindings [Name] Expr) Expr
    -- Access
    | Car Expr
    | Cdr Expr
    | VectorRef Expr Int
    -- Construction
    | Cons Expr Expr
    | Vector [Expr]
    | Values [Expr]
    | Lambda Name Expr
    -- Procedure call
    | ProcCall Name [Expr]
      deriving (Eq, Show, Typeable, Data)

type ShapedName = (Name, Shape)

newtype Bindings p v = Bindings [(p, v)] deriving (Eq, Show, Typeable, Data)

instance Functor (Bindings p) where
    fmap f (Bindings bs)
        = Bindings [(x, f v) | (x, v) <- bs]

instance Foldable (Bindings p) where
    foldMap f (Bindings bs)
        = mconcat [f v | (_, v) <- bs]

instance Traversable (Bindings p) where
    traverse f (Bindings bs)
        = Bindings <$> sequenceA [liftA2 (,) (pure x) (f v) | (x, v) <- bs]

mkLet :: [(Name, Expr)] -> Expr -> Expr
mkLet [] body = body
mkLet bs body = Let (Bindings bs) body

mkLetValues :: [([Name], Expr)] -> Expr -> Expr
mkLetValues [] body = body
mkLetValues bs body = LetValues (Bindings bs) body

smartLetValues :: [([Name], Expr)] -> Expr -> Expr
smartLetValues bs body = mkLet bs1 (mkLetValues bs2 body)
    where
      bs1 = [(x,  e) | ([x],        e) <- bs]
      bs2 = [(xs, e) | (xs@(_:_:_), e) <- bs]

procName :: Defn -> Name
procName (Defn (proc_name, _) _ _) = proc_name

-- Free variables
fv :: Expr -> [Name]
fv (Var x)    = [x]
fv Nil        = []
fv (Bool _)   = []
fv (Real _)   = []
fv (If p c a) = (fv p) ++ (fv c) ++ (fv a)
fv (Let (Bindings bs) body) = (fv body ++ concatMap fv es) \\ xs
    where
      (xs, es) = unzip bs
fv (LetValues (Bindings bs) body) = (fv body ++ concatMap fv es) \\ concat xss
    where
      (xss, es) = unzip bs
fv (Car e)           = fv e
fv (Cdr e)           = fv e
fv (VectorRef e _)   = fv e
fv (Cons e1 e2)      = fv e1 ++ fv e2
fv (Vector es)       = concatMap fv es
fv (Values es)       = concatMap fv es
fv (Lambda x e)      = delete x (fv e)
fv (ProcCall _ args) = concatMap fv args

-- Pretty-printing of programs

instance Pretty Prog where
    pp (Prog []    expr) = pp expr
    pp (Prog defns expr) = parens $ symbol "begin" $+$ body
        where
          body = nest 1 (sepMap pp defns $+$ pp expr)

instance Pretty Defn where
    pp (Defn proc args body) = parens $ proto $+$ nest 1 (shapes $+$ pp body)
        where
          (proc_name, proc_shape) = proc
          (arg_names, arg_shapes) = unzip args
          proto  =     symbol "define"
                   <+> parens (pp proc_name <+> sepMap pp arg_names)
          shapes = ppForm "argument-types" (arg_shapes ++ [proc_shape])

instance Pretty Shape where
    pp NilSh          = ppList []
    pp RealSh         = symbol "real"
    pp BoolSh         = symbol "bool"
    pp (ConsSh t1 t2) = ppForm "cons"   [t1, t2]
    pp (VectorSh ts)  = ppForm "vector" ts
    pp (ValuesSh ts)  = ppForm "values" ts
    pp FunctionSh     = symbol "escaping-function"

instance Pretty Expr where
    pp Nil          = ppList []
    pp (Var x)      = pp x
    pp (Real r)     = real r
    pp (Bool True)  = symbol "#t"
    pp (Bool False) = symbol "#f"

    pp (If p c a) = parens $ symbol "if" <+> sepMap pp [p, c, a]

    pp (Let (Bindings bs) e)
        = parens $ (symbol "let" <+> bindings) $+$ (nest 1 body)
        where
          body = pp e
          bindings = parens . sepMap ppBinding $ bs
          ppBinding (name, expr) = ppList [pp name, pp expr]

    pp (LetValues (Bindings bs) e)
        = parens $ (symbol "let-values" <+> bindings) $+$ (nest 1 body)
        where
          body = pp e
          bindings = parens . sepMap ppBinding $ bs
          ppBinding (names, expr) = ppList [ppList (map pp names), pp expr]

    pp (Car e)         = ppForm "car" [e]
    pp (Cdr e)         = ppForm "cdr" [e]
    pp (VectorRef e i) = ppList [symbol "vector-ref", pp e, int i]
    pp (Cons e1 e2)    = ppForm "cons"   [e1, e2]
    pp (Vector es)     = ppForm "vector" es
    pp (Values es)     = ppForm "values" es
    pp (Lambda x e)    = ppList [symbol "lambda", parens (pp x), pp e]

    pp (ProcCall proc args) = ppList $ pp proc : map pp args

-- Annotated programs, definitions, and expressions

-- It is really obnoxious that I have to write down these definitions
-- and pretty trivial conversion functions but there appears to be no
-- better way to arrange these things.  Cf.
--
-- Simon L. Peyton Jones, David Lester,
-- /A Modular Fully-lazy Lambda Lifter in Haskell/,
-- Software --- Practice and Experience, Vol. 21(5), 479--506 (1991).

type AnnProg ann = (ann, AnnProg' ann)
data AnnProg' ann
    = AnnProg [AnnDefn ann] (AnnExpr ann)
      deriving (Eq, Show)

stripAnnProg :: AnnProg ann -> Prog
stripAnnProg (_, prog) = stripAnnProg' prog

stripAnnProg' :: AnnProg' ann -> Prog
stripAnnProg' (AnnProg defns expr)
    = Prog (map stripAnnDefn defns) (stripAnnExpr expr)

instance Pretty (AnnProg' ann) where
    pp = pp . stripAnnProg'

type AnnDefn ann = (ann, AnnDefn' ann)
data AnnDefn' ann
    = AnnDefn ShapedName [ShapedName] (AnnExpr ann)
      deriving (Eq, Show)

stripAnnDefn :: AnnDefn ann -> Defn
stripAnnDefn (_, defn) = stripAnnDefn' defn

stripAnnDefn' :: AnnDefn' ann -> Defn
stripAnnDefn' (AnnDefn proc args body)
    = Defn proc args (stripAnnExpr body)

instance Pretty (AnnDefn' ann) where
    pp = pp . stripAnnDefn'

type AnnExpr ann = (ann, AnnExpr' ann)
data AnnExpr' ann
    = AnnVar Name
    | AnnNil
    | AnnBool Bool
    | AnnReal Real
    | AnnIf (AnnExpr ann) (AnnExpr ann) (AnnExpr ann)
    | AnnLet (Bindings Name (AnnExpr ann)) (AnnExpr ann)
    | AnnLetValues (Bindings [Name] (AnnExpr ann)) (AnnExpr ann)
    | AnnCar (AnnExpr ann)
    | AnnCdr (AnnExpr ann)
    | AnnVectorRef (AnnExpr ann) Int
    | AnnCons (AnnExpr ann) (AnnExpr ann)
    | AnnVector [AnnExpr ann]
    | AnnValues [AnnExpr ann]
    | AnnLambda Name (AnnExpr ann)
    | AnnProcCall Name [AnnExpr ann]
      deriving (Eq, Show, Typeable, Data)

stripAnnExpr :: AnnExpr ann -> Expr
stripAnnExpr (_, e) = stripAnnExpr' e

stripAnnExpr' :: AnnExpr' ann -> Expr
stripAnnExpr' (AnnVar x)  = Var x
stripAnnExpr' AnnNil      = Nil
stripAnnExpr' (AnnBool b) = Bool b
stripAnnExpr' (AnnReal r) = Real r
stripAnnExpr' (AnnIf p c a)
    = If (stripAnnExpr p) (stripAnnExpr c) (stripAnnExpr a)
stripAnnExpr' (AnnLet bindings body)
    = Let (fmap stripAnnExpr bindings) (stripAnnExpr body)
stripAnnExpr' (AnnLetValues bindings body)
    = LetValues (fmap stripAnnExpr bindings) (stripAnnExpr body)
stripAnnExpr' (AnnCar e)              = Car (stripAnnExpr e)
stripAnnExpr' (AnnCdr e)              = Cdr (stripAnnExpr e)
stripAnnExpr' (AnnVectorRef e i)      = VectorRef (stripAnnExpr e) i
stripAnnExpr' (AnnCons e1 e2)         = Cons (stripAnnExpr e1)
                                             (stripAnnExpr e2)
stripAnnExpr' (AnnVector es)          = Vector (map stripAnnExpr es)
stripAnnExpr' (AnnValues es)          = Values (map stripAnnExpr es)
stripAnnExpr' (AnnLambda x e)         = Lambda x (stripAnnExpr e)
stripAnnExpr' (AnnProcCall proc args) = ProcCall proc (map stripAnnExpr args)

instance Pretty (AnnExpr' ann) where
    pp = pp . stripAnnExpr'

-- Annotated shapes

data AnnShape ann
    = AnnNilSh  ann
    | AnnRealSh ann
    | AnnBoolSh ann
    | AnnConsSh (AnnShape ann) (AnnShape ann)
    | AnnVectorSh [AnnShape ann]
    | AnnValuesSh [AnnShape ann]
    | AnnFunctionSh ann
      deriving (Eq, Show)

stripAnnShape :: AnnShape ann -> Shape
stripAnnShape (AnnNilSh  _)     = NilSh
stripAnnShape (AnnRealSh _)     = RealSh
stripAnnShape (AnnBoolSh _)     = BoolSh
stripAnnShape (AnnConsSh s1 s2) = ConsSh (stripAnnShape s1) (stripAnnShape s2)
stripAnnShape (AnnVectorSh ss)  = VectorSh (map stripAnnShape ss)
stripAnnShape (AnnValuesSh ss)  = ValuesSh (map stripAnnShape ss)
stripAnnShape (AnnFunctionSh _) = FunctionSh

fringe :: AnnShape ann -> [(ann, Shape)]
fringe (AnnNilSh  ann)   = [(ann, NilSh )]
fringe (AnnRealSh ann)   = [(ann, RealSh)]
fringe (AnnBoolSh ann)   = [(ann, BoolSh)]
fringe (AnnConsSh s1 s2) = fringe s1 ++ fringe s2
fringe (AnnVectorSh ss)  = concatMap fringe ss
fringe (AnnValuesSh ss)  = concatMap fringe ss
fringe (AnnFunctionSh ann) = [(ann, FunctionSh)]

annots :: AnnShape ann -> [ann]
annots = map fst . fringe
