-- Needed in general
{-# LANGUAGE GADTs, KindSignatures, RankNTypes #-}

{-# OPTIONS_GHC -Wall -fwarn-tabs #-}

----------------------------------------------------------------
--                                                    2011.06.07
-- |
-- Module      :  HOAS_simple
-- Copyright   :  Copyright (c) 2009 wren ng thornton
-- License     :  BSD
-- Maintainer  :  wren@community.haskell.org
-- Stability   :  experimental
-- Portability :  non-portable (GADTs,...)
--
-- This module provides a type-safe HOAS GADT implementation
-- of lambda calculi, with conversion to the de Bruijn GADT, ala:
--
--    * http://www.cse.unsw.edu.au/~chak/haskell/term-conv/
--
----------------------------------------------------------------
module HOAS_simple
    (
    -- * Object language terms
      HTerm(..)
    , Term(..)
    , eval
    -- * Conversion from HOAS to de Bruijn
    , convert
    , convert'
    ) where

import Data.Typeable
import DeBruijn_simple (Ix(..), DTerm(..))

----------------------------------------------------------------
----------------------------------------------------------------
-- | Abstract terms are constants, applications, and lambda
-- abstractions. The @Term@ class is for abstracting over 'HTerm'
-- and 'CTerm' in order to: (1) not allow the Haskell function
-- passed to @HLam@ perform case analysis on its argument, and (2)
-- to only have the @CVar@ constructor present during conversion.
-- Thus we avoid so-called \"exotic terms\" which do not correspond
-- to anything in lambda calculus.
--
-- All the types must be 'Typeable' so that we can perform a 'gcast'
-- when converting to de Bruijn format. Unfortunately this means
-- this isn't simply an extension of 'Applicative' which adds @lam@.
class Term t where
    con :: (Typeable a)             => a -> t a
    app :: (Typeable a, Typeable b) => t (a -> b) -> t a -> t b
    lam :: (Typeable a, Typeable b) => (t a -> t b) -> t (a -> b)

-- same as (<$>) and (<*>)
infixl 4 `app`


-- | Concrete HOAS terms. We have no variables because this is HOAS,
-- which means we use the host language's binding constructs to
-- provide object language variable names.
data HTerm :: * -> * where
    HCon   :: (Typeable a)            => a -> HTerm a
    HApp   :: (Typeable a,Typeable b) => HTerm (a -> b) -> HTerm a -> HTerm b
    HLam   :: (Typeable a,Typeable b) => (HTerm a -> HTerm b) -> HTerm (a -> b)
    
    -- Only used for conversion from de Bruijn
    HVar   :: (Typeable a) => Integer -> HTerm a

instance Term HTerm where
    con = HCon
    app = HApp
    lam = HLam

-- | Evaluate a HOAS term.
eval :: HTerm a -> a
eval (HCon c)     = c
eval (HApp e1 e2) = eval e1 $ eval e2
eval (HLam f)     = eval . f . HCon
eval (HVar _n)    = error "Can't use HVar outside of conversion"


----------------------------------------------------------------
----------------------------------------------------------------

-- | The layout of a binding environment (or typing context) is a
-- stack containing all the valid indices for that binding environment.
-- The first parameter captures the typing environment of the 'Env'
-- and threads it through all the 'Ix', the second parameter captures
-- the typing environment of the layout itself. They are unified
-- in the 'convert' function.
data Layout :: * -> * -> * where
    EmptyLayout :: Layout r ()
    PushLayout  :: (Typeable a) => Layout r r' -> Ix r a -> Layout r (r',a)


-- | /O(n)/. The number of entries in an environment layout, which
-- is \"coincidentally\" the same as the next free variable.
size :: Layout r r' -> Integer
size = go 0
    where
    go :: Integer -> Layout r r' -> Integer
    go i EmptyLayout        = i
    go i (PushLayout lyt _) = (go $! i+1) lyt


-- | /O(n)/. Add a new entry to the layout, incrementing all the
-- 'Ix' which have already been allocated.
extendLayout :: (Typeable a) => Layout r r' -> Layout (r,a) (r',a)
extendLayout = \lyt -> PushLayout (fmapSucc lyt) ZeroIx
    where
    fmapSucc :: Layout r r' -> Layout (r,a) r'
    fmapSucc EmptyLayout         = EmptyLayout
    fmapSucc (PushLayout lyt ix) = PushLayout (fmapSucc lyt) (SuccIx ix)


-- | Project the nth index out of an environment layout.
lookupLayout :: (Typeable a) => Integer -> Layout r r' -> Ix r a
lookupLayout = start
    where
    start n
        | n < 0     = error "HOAS.lookupLayout: negative integer"
        | otherwise = go n
    
    go :: (Typeable a) => Integer -> Layout r r' -> Ix r a
    go _ EmptyLayout        = error "HOAS.lookupLayout: internal error"
    go 0 (PushLayout _  ix) =
        case gcast ix of -- alas, this seems necessary.
        Nothing  -> error "HOAS.lookupLayout: gcast error"
        Just ix' -> ix'
    go n (PushLayout lyt _) = (lookupLayout $! n-1) lyt


-- | Convert a (closed) HOAS term into a closed de Bruijn term. We
-- hand out tags in the order that we encounter bindings, hence they
-- are the inverse of the de Bruijn indices.
convert :: HTerm a -> DTerm () a
convert = go EmptyLayout
    where
    go :: Layout r r -> HTerm a -> DTerm r a
    go lyt (HVar n)     = DVar (lookupLayout (size lyt - n - 1) lyt)
    go _   (HCon c)     = DCon c
    go lyt (HApp e1 e2) = DApp (go lyt e1) (go lyt e2)
    go lyt (HLam f)     = DLam (go (extendLayout lyt) (f (HVar $ size lyt)))


-- | A variant of 'convert' which memoizes the layout's size. This
-- should be significantly faster since we save /O(n)/ for every
-- @HLam@ and every time the @HLam@'s function uses its argument,
-- where /n/ is the size of the environment at that point.
convert' :: HTerm a -> DTerm () a
convert' = go 0 EmptyLayout
    where
    go :: Integer -> Layout r r -> HTerm a -> DTerm r a
    go sz lyt (HVar n)     = DVar $ lookupLayout (sz - n - 1) lyt
    go _  _   (HCon c)     = DCon c
    go sz lyt (HApp e1 e2) = DApp (go sz lyt e1) (go sz lyt e2)
    go sz lyt (HLam f)     =
        DLam $ (go $! sz+1) (extendLayout lyt) (f (HVar sz))

----------------------------------------------------------------
----------------------------------------------------------- fin.