{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
-- | Expression for /lambda abstraction/s
-- in /Higher-Order Abstract Syntax/ (HOAS).
module Language.Symantic.Expr.Lambda
 ( module Language.Symantic.Expr.Lambda
 , Sym_Lambda_Lam(..)
 ) where

import qualified Control.Applicative as Applicative
import qualified Data.Function as Fun
import Data.Monoid
import Data.Proxy (Proxy(..))
import Data.Text (Text)
import qualified Data.Text as Text
import Data.Type.Equality ((:~:)(Refl))
import Prelude hiding (const, id)

import Language.Symantic.Type
import Language.Symantic.Repr
import Language.Symantic.Expr.Root
import Language.Symantic.Expr.Error
import Language.Symantic.Expr.From
import Language.Symantic.Trans.Common

-- * Class 'Sym_Lambda'
-- | Symantic.
class Sym_Lambda_Lam repr => Sym_Lambda repr where
	-- | /Lambda application/.
	($$) :: repr ((->) arg res) -> repr arg -> repr res
	default ($$) :: Trans t repr
	 => t repr ((->) arg res) -> t repr arg -> t repr res
	($$) f x = trans_lift (trans_apply f $$ trans_apply x)
	
	-- | Convenient 'lam' and '$$' wrapper.
	let_ :: repr var -> (repr var -> repr res) -> repr res
	let_ x y = lam y $$ x
	
	id :: repr a -> repr a
	id a = (lam Fun.id) $$ a
	
	const :: repr a -> repr b -> repr a
	const a b = lam (lam . Fun.const) $$ a $$ b
	
	-- | /Lambda composition/.
	(#) :: repr (b -> c) -> repr (a -> b) -> repr (a -> c)
	(#) f g = lam $ \a -> f $$ (g $$ a)
	
	flip :: repr (a -> b -> c) -> repr (b -> a -> c)
	flip f = lam $ \b -> lam $ \a -> f $$ a $$ b

infixl 0 $$
infixr 9 #

instance Sym_Lambda Repr_Host where
	($$)  = (Applicative.<*>)
instance Sym_Lambda Repr_Text where
	-- ($$) = repr_text_infix "$" (Precedence 0)
	($$) (Repr_Text a1) (Repr_Text a2) =
		Repr_Text $ \p v ->
			let p' = precedence_App in
			paren p p' $ a1 p' v <> " " <> a2 p' v
	let_ e in_ =
		Repr_Text $ \p v ->
			let p' = Precedence 2 in
			let x = "x" <> Text.pack (show v) in
			paren p p' $ "let" <> " " <> x <> " = "
			 <> unRepr_Text e (Precedence 0) (succ v) <> " in "
			 <> unRepr_Text (in_ (Repr_Text $ \_p _v -> x)) p' (succ v)
	(#)   = repr_text_infix "." (Precedence 9)
	id    = repr_text_app1 "id"
	const = repr_text_app2 "const"
	flip  = repr_text_app1 "flip"
instance (Sym_Lambda r1, Sym_Lambda r2) => Sym_Lambda (Repr_Dup r1 r2) where
	($$) = repr_dup2 sym_Lambda ($$)

sym_Lambda :: Proxy Sym_Lambda
sym_Lambda = Proxy

-- * Type 'Expr_Lambda'
-- | Expression.
data Expr_Lambda (root:: *)
type instance Root_of_Expr      (Expr_Lambda root)      = root
type instance Type_of_Expr      (Expr_Lambda root)      = Type_Fun
type instance Sym_of_Expr       (Expr_Lambda root) repr = Sym_Lambda repr
type instance Error_of_Expr ast (Expr_Lambda root)      = Error_Expr_Lambda ast

-- | Parsing utility to check that the given type is a 'Type_Fun'
-- or raise 'Error_Expr_Type_mismatch'.
check_type_fun
 :: forall ast ex root ty h ret.
 ( root ~ Root_of_Expr ex
 , ty ~ Type_Root_of_Expr ex
 , Type0_Lift   Type_Fun (Type_of_Expr root)
 , Type0_Unlift Type_Fun (Type_of_Expr root)
 , Error_Expr_Lift (Error_Expr (Error_of_Type ast ty) ty ast)
                   (Error_of_Expr ast root)
 )
 => Proxy ex -> ast -> ty h
 -> (Type_Fun ty h -> Either (Error_of_Expr ast root) ret)
 -> Either (Error_of_Expr ast root) ret
check_type_fun ex ast ty k =
	case type0_unlift $ unType_Root ty of
	 Just ty_f -> k ty_f
	 Nothing -> Left $ error_expr ex $
		Error_Expr_Type_mismatch ast
		 (Exists_Type0 (type_var0 SZero `type_fun` type_var0 (SSucc SZero)
		 :: ty ((->) Var0 Var0)))
		 (Exists_Type0 ty)

-- | Parse a /lambda variable/.
var_from
 :: forall ast root hs ret.
 ( Type0_From ast (Type_Root_of_Expr root)
 , Error_Expr_Lift (Error_Expr_Lambda ast)
                   (Error_of_Expr ast root)
 , Root_of_Expr root ~ root
 ) => Text -> ExprFrom ast (Expr_Lambda root) hs ret
var_from name _ex ast = go
	where
	go :: forall ex hs'. (ex ~ (Expr_Lambda root))
	 => Lambda_Context (Lambda_Var (Type_Root_of_Expr ex)) hs'
	 -> ( forall h. Type_Root_of_Expr ex h
	    -> Forall_Repr_with_Context ex hs' h
	    -> Either (Error_of_Expr ast (Root_of_Expr ex)) ret )
	 ->    Either (Error_of_Expr ast (Root_of_Expr ex)) ret
	go c k' =
		case c of
		 Lambda_Context_Empty -> Left $ error_expr_lift $
			Error_Expr_Lambda_Var_unbound name ast
		 Lambda_Var n ty `Lambda_Context_Next` _ | n == name ->
			k' ty $ Forall_Repr_with_Context $
			 \(repr `Lambda_Context_Next` _) -> repr
		 _ `Lambda_Context_Next` ctx' ->
			go ctx' $ \ty (Forall_Repr_with_Context repr) ->
			k' ty $ Forall_Repr_with_Context $
			 \(_ `Lambda_Context_Next` c') -> repr c'

-- | Parse '$$'.
app_from
 :: forall ty ast root hs ret.
 ( ty ~ Type_Root_of_Expr root
 , Type0_From ast ty
 , Type0_Eq ty
 , Expr_From ast root
 , Type0_Lift Type_Fun (Type_of_Expr root)
 , Error_Expr_Lift (Error_Expr (Error_of_Type ast ty) ty ast)
                   (Error_of_Expr ast root)
 , Type0_Unlift Type_Fun (Type_of_Expr root)
 , Root_of_Expr root ~ root
 ) => ast -> ast
 -> ExprFrom ast (Expr_Lambda root) hs ret
app_from ast_lam ast_arg_actual ex ast ctx k =
	expr_from (Proxy::Proxy root) ast_lam ctx $
	 \(ty_lam::ty h_lam) (Forall_Repr_with_Context l) ->
	expr_from (Proxy::Proxy root) ast_arg_actual ctx $
	 \ty_arg_actual (Forall_Repr_with_Context arg_actual) ->
	case type0_unlift $ unType_Root ty_lam of
	 Nothing -> Left $ error_expr ex $
		Error_Expr_Type_mismatch ast
		 (Exists_Type0 (type_var0 SZero `type_fun` type_var0 (SSucc SZero)
		 :: ty ((->) Var0 Var0)))
		 (Exists_Type0 ty_lam)
	 Just (Type2 Proxy ty_arg_expected ty_res
	  :: Type_Fun ty h_lam) ->
		check_type0_eq ex ast ty_arg_expected ty_arg_actual $ \Refl ->
		k ty_res $ Forall_Repr_with_Context $
		 \c -> l c $$ arg_actual c

-- | Parse 'lam'.
lam_from
 :: forall ty ast root hs ret.
 ( ty ~ Type_Root_of_Expr root
 , root ~ Root_of_Expr root
 , Type0_From ast ty
 , Expr_From ast root
 , Type0_Lift Type_Fun (Type_of_Expr root)
 , Error_Expr_Lift (Error_Expr (Error_of_Type ast ty) ty ast)
                   (Error_of_Expr ast root)
 ) => Text -> ast -> ast
 -> ExprFrom ast (Expr_Lambda root) hs ret
lam_from name ast_ty_arg ast_body ex ast ctx k =
	either (\err -> Left $ error_expr ex $ Error_Expr_Type err ast) Fun.id $
	type0_from (Proxy::Proxy ty) ast_ty_arg $ \ty_arg -> Right $
	expr_from (Proxy::Proxy root) ast_body
	 (Lambda_Var name ty_arg `Lambda_Context_Next` ctx) $
	 \ty_res (Forall_Repr_with_Context res) ->
	k (ty_arg `type_fun` ty_res) $ Forall_Repr_with_Context $
	 \c -> lam $ \arg -> res (arg `Lambda_Context_Next` c)

-- | Parse 'let_'.
let_from
 :: forall ty ast root hs ret.
 ( ty ~ Type_Root_of_Expr root
 , root ~ Root_of_Expr root
 , Type0_From ast ty
 , Expr_From ast root
 , Error_Expr_Lift (Error_Expr (Error_of_Type ast ty) ty ast)
                   (Error_of_Expr ast root)
 ) => Text -> ast -> ast
 -> ExprFrom ast (Expr_Lambda root) hs ret
let_from name ast_var ast_body _ex _ast ctx k =
	expr_from (Proxy::Proxy root) ast_var ctx $
	 \ty_var (Forall_Repr_with_Context var) ->
	expr_from (Proxy::Proxy root) ast_body
	 (Lambda_Var name ty_var `Lambda_Context_Next` ctx) $
	 \ty_res (Forall_Repr_with_Context res) ->
	k ty_res $ Forall_Repr_with_Context $
	 \c -> let_ (var c) $ \arg -> res (arg `Lambda_Context_Next` c)

-- * Type 'Error_Expr_Lambda'
data Error_Expr_Lambda ast
 =   Error_Expr_Lambda_Var_unbound Lambda_Var_Name ast
 deriving (Eq, Show)