"flat" circular substitution

FlatCircularSubstitution.hs

{-# language DeriveTraversable #-}
{-# language BlockArguments #-}
{-# language LambdaCase #-}
{-# language PatternSynonyms #-}
{-# language ViewPatterns #-}

import Control.Monad (ap)
import Control.Monad.ST
import Control.Monad.Fix
import Data.Function (on)
import Data.Functor ((<&>))
import Data.Primitive.Array
import Data.STRef
import Data.Void
import Prelude hiding (pi)
import Control.Monad.ST.Unsafe
import Text.Read

type Level = Int
data Name = Name { hint :: String, nameId :: Occ } -- lazy nameId

instance Show Name where
  showsPrec d (Name n o) = showParen (d >= 11) $ showString "Name " . showsPrec 11 n . showChar ' ' . showsPrec 11 o

instance Read Name where
  readPrec  = parens $ prec 10 do
    Ident "Name" <- lexP
    Name <$> step readPrec <*> step readPrec

type Occ = Int

instance Eq Name where 
  (==) = (==) `on` nameId

instance Ord Name where
  compare = compare `on` nameId

data Constant
  = Level
  | LevelLiteral
  | Omega
  deriving (Eq,Ord,Show,Read)

data Term a
  = Free a
  | Bound Occ
  | Constant !Constant
  | Term a :+ {-# unpack #-} !Level
  | Max [Term a]
  | Type !(Term a)
  | LAM {-# unpack #-} !Name !(Term a) !(Term a)
  | PI {-# unpack #-} !Name !(Term a) !(Term a)
  | SIGMA {-# unpack #-} !Name !(Term a) !(Term a)
  | App !(Term a) !(Term a)
  | Fst !(Term a)
  | Snd !(Term a)
  | Pair !(Term a) !(Term a) !(Term a)
  deriving (Show, Read, Eq, Ord, Functor, Foldable, Traversable)

class Bindable t where
  bound :: t -> Occ
  bind :: Occ -> t

prime :: Bindable t => t -> Occ
prime t = bound t + 1

instance Bindable (Term a) where
  bound Free{} = 0
  bound Bound{} = 0
  bound Constant{} = 0
  bound (a :+ _) = bound a
  bound (Max xs) = foldr (\a r -> bound a `max` r) 0 xs
  bound (Type t) = bound t
  bound (LAM b t _) = nameId b `max` bound t
  bound (PI b t _) = nameId b `max` bound t
  bound (SIGMA b t _) = nameId b `max` bound t
  bound (App x y) = bound x `max` bound y
  bound (Fst t) = bound t
  bound (Snd t) = bound t
  bound (Pair t x y) = bound t `max` bound x `max` bound y
  bind = Bound

binder
  :: Bindable t
  => (Name -> t -> r)
  -> String
  -> (t -> t)
  -> r
binder c n e = c (Name n i) body where
  body = e (bind i)
  i = prime body

binderST
  :: Bindable t
  => (Name -> t -> r)
  -> String
  -> (t -> ST s t)
  -> ST s r
binderST c n e = do
  ix <- newSTRef (error "binderM: fix failed") -- mfix with a tuple would generalize to any MonadFix
  body <- mfix \body -> do
    let i = prime body
    writeSTRef ix i
    e (bind i)
  readSTRef ix <&> \i -> c (Name n i) body

lamST, piST, sigmaST :: String -> Term a -> (Term a -> ST s (Term a)) -> ST s (Term a)
lamST s t = binderST (`LAM` t) s
piST s t = binderST (`PI` t) s
sigmaST s t = binderST (`SIGMA` t) s

rebind :: MutableArray s (Term b) -> Term a -> (a -> Term b) -> ST s (Term b)
rebind env xs0 f = go xs0 where
  instantiate = writeArray env . nameId
  go = \case
    Free a -> pure $ f a
    Bound b -> readArray env b
    Constant c -> pure $ Constant c
    m :+ n -> (:+ n) <$> go m
    Type t -> Type <$> go t
    Max xs -> Max <$> traverse go xs
    LAM b t e -> do
      t' <- go t 
      lamST (hint b) t' \v -> instantiate b v *> go e
    PI b t e -> do
      t' <- go t
      piST (hint b) t' \v -> instantiate b v *> go e
    SIGMA b t e -> do
      t' <- go t
      sigmaST (hint b) t' \v -> instantiate b v *> go e
    App x y -> App <$> go x <*> go y
    Fst x -> Fst <$> go x
    Snd x -> Snd <$> go x
    Pair t x y -> Pair <$> go t <*> go x <*> go y

instance Applicative Term where
  pure = Free
  (<*>) = ap

instance Monad Term where
  return = Free
  m >>= f = runST do
    env <- newArray (prime m) $ error "(Term.>>=): unbound variable"
    rebind env m f

bind1 :: Occ -> Term a -> Term a -> Term a
bind1 i e v = runST do
  env <- newArray (i + 1) $ error "bind1: unbound variable"
  writeArray env i v
  rebind env v pure

pattern Lam :: String -> Term a -> (Term a -> Term a) -> Term a
pattern Lam s t f <- LAM (Name s i) t (bind1 i -> f) where
  Lam s t f = binder (`LAM` t) s f

pattern Pi :: String -> Term a -> (Term a -> Term a) -> Term a
pattern Pi s t f <- PI (Name s i) t (bind1 i -> f) where
  Pi s t f = binder (`PI` t) s f

pattern Sigma :: String -> Term a -> (Term a -> Term a) -> Term a
pattern Sigma s t f <- SIGMA (Name s i) t (bind1 i -> f) where
  Sigma s t f = binder (`SIGMA` t) s f

{-# complete Free, Bound, Constant, (:+), Type, Max, Lam, Pi, Sigma, App, Fst, Snd, Pair #-}

subst :: Eq a => a -> Term a -> Term a -> Term a
subst a x y = y >>= \a' -> if a == a' then x else Free a'

main :: IO ()
main = do
  let 
    id_ :: Term Void
    id_ = Lam "x" (Constant Omega) \x -> x 
    const_ :: Term Void
    const_ = Lam "x" (Constant Omega) \x -> Lam "y" (Constant Omega) \y -> x
    ki :: Term String
    ki = App (Free "const") (Free "id")
  print id_
  print const_
  print $ (Lam "_" (Constant Omega) \_ -> ki) >>= \case
    "id" -> id_
    "const" -> const_

Based on http://comonad.com/reader/2014/fast-circular-substitution/

The superpower you are given by this notion of names is that you get to know how deep the binding structure of any syntax tree goes, and at least if you maintain the invariant of always using the HOAS-like view patterns to work with terms, these names are densely packed.

Modified to use Data.Primitive.Array instead of an IntMap and to provide view patterns that offer a fully faux HOAS API.