{-# LANGUAGE DeriveFunctor, GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleContexts, LambdaCase #-}

-- | A humane representation of binding.
module ABT ( FreeVar (..)
           , Operator (..)
           , ABT
           , View (..)
           , view
           , fold
           , subst ) where
import Control.Monad.Gen
import Control.Applicative

data Operator = Operator { operArity :: Int
                         , operName  :: String }
              deriving(Eq)

newtype FreeVar = Free Int
                deriving(Eq, Ord, Enum)

-- | Internal representation of an ABT. This is locally nameless and
-- somewhat painful to work with. Happily with this library we write
-- this code once and rely on a nicer interface in the future. This
-- does easily admit alpha-equivalence though.
data ABT = FreeVar FreeVar
         | BoundVar Int
         | BindIn ABT
         | App Operator [ABT]
         deriving (Eq)

-- | The humane view of an ABT.
data View t = Var FreeVar -- ^ A variable
            | Binding FreeVar t -- ^ A binding. Binds a FreeVar in t
            | Oper Operator [t] -- ^ Applying some operator to some ts.
            deriving Functor

-- | Given a freevar, replace the outmost DB variable with it.
unbind :: FreeVar -> ABT -> ABT
unbind var = go 0
  where go i = \case
          BoundVar j -> if i == j then FreeVar var else BoundVar j
          FreeVar v -> FreeVar v
          BindIn a -> BindIn (go (i + 1) a)
          App oper abts -> App oper (map (go i) abts)

-- | Add a new binder which binds a particular free variable. This
-- replaces all occurences of this var with a DB var pointing to this
-- outermost binder.
bind :: FreeVar -> ABT -> ABT
bind var = go 0
  where go i = \case
          FreeVar var' -> if var == var' then BoundVar i else FreeVar var'
          BoundVar j -> BoundVar j
          BindIn a -> BindIn (go (i + 1) a)
          App oper abts -> App oper (map (go i) abts)

-- | Given an ABT unfold it partially and give a view of the
-- structure. This can potentially create new fresh variables so we
-- need to have effects here.
view :: (Functor m, MonadGen FreeVar m) => ABT -> m (View ABT)
view = \case
  FreeVar v -> return $ Var v
  BoundVar _ -> error "Impossible"
  App oper abts -> return $ Oper oper abts
  BindIn a -> (Binding <*> flip unbind a) <$> gen

-- | Given a view, we can smush it back into the efficient, locally
-- nameless representation.
fold :: View ABT -> ABT
fold = \case
  Var v -> FreeVar v
  Binding v a -> bind v a
  Oper oper abts -> App oper abts

-- | A priveleged subst. This way avoids unneeded calls to gen.
subst :: ABT -> FreeVar -> ABT -> ABT
subst new var = go
  where go = \case
          FreeVar var' -> if var == var' then new else FreeVar var'
          BoundVar i -> BoundVar i
          BindIn a -> BindIn (go a)
          App oper abts -> App oper (map go abts)