{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeSynonymInstances  #-}

module FoldTest where

import Base
import Fold
import AST2
import List

import qualified Data.Map as M
import Data.Maybe

-- Length of lists

data K' a b = K' {unK' :: a}

len_alg :: Alg (List a) (K' Int)
len_alg _ (L (K ())) = K' 0
len_alg _ (R (K x :*: Id (K' xs))) = K' (1 + xs)

len_alg' :: Alg (List a) (K' Int)
len_alg' _ = (const (K' 0)) 
           & (fn (\ (K x) (Id (K' xs)) -> K' (1 + xs)))

len :: [a] -> Int
len = unK' . fold List len_alg

len' :: [a] -> Int
len' = unK' . fold List len_alg'

-- Evaluation of abstract syntax trees

type Env = M.Map Var Int

data family Value a :: *
data instance Value Expr = EV (Env -> Int)
data instance Value Decl = DV (Env -> Env)
data instance Value Var  = VV Var

-- A Hack:
infixr &|
(&|) x y = tag x & y

-- The algebra:
eval_alg :: Alg AST Value
eval_alg _ =  (\ (K x)                       -> EV (const x))
           &| (fn (\ (Id (EV x)) (Id (EV y)) -> EV (\ env -> x env + y env)))
           &| (fn (\ (Id (EV x)) (Id (EV y)) -> EV (\ env -> x env * y env)))
           &| (\ (Id (VV x))                 -> EV (fromJust . M.lookup x))
           &| (fn (\ (Id (DV e)) (Id (EV x)) -> EV (\ env -> x (e env))))
           &| (fn (\ (Id (VV x)) (Id (EV v)) -> DV (\ env -> M.insert x (v env) env )))
           &| (fn (\ (Id (DV e)) (Id (DV f)) -> DV (f . e)))
           &| tag (\ (K x)                   -> VV x)

{-
 - Ideal way to write the algebra:
 -

eval_alg :: Alg AST Value
eval_alg _ =  (\ x                           -> EV (const x))
           &| (\ (EV x) (EV y)               -> EV (\ env -> x env + y env)))
           &| (\ (EV x) (EV y)               -> EV (\ env -> x env * y env)))
           &| (\ (VV x)                      -> EV (fromJust . M.lookup x))
           &| (\ (DV e) (EV x)               -> EV (\ env -> x (e env))))
           &| (\ (VV x) (EV v)               -> DV (\ env -> M.insert x (v env) env )))
           &| (\ (DV e) (DV f)               -> DV (f . e)))
           &| (\ (K x)                       -> VV x)
-}


eval :: Expr -> Int
eval = (\ (EV x) -> x M.empty) . fold Expr eval_alg 

expr :: Expr
expr = Let ("x" := Add (Const 2) (Const 3)) (Let ("y" := Add (EVar "x") (Const 1)) (Mul (EVar "x") (EVar "y")))