Hash consing

intern.hs

{-# LANGUAGE CPP #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleInstances #-}
module Intern where

import GHC.Generics (Generic)
import Data.Function (on)

import Data.Hashable
import Data.Interned

#define _UPK(_x) {-# UNPACK #-} !(_x)

type Var = Int

data Term
  = App _UPK(Id) !Term !Term
  | Lam _UPK(Id) _UPK(Var) !Term !Term
  | Pi  _UPK(Id) _UPK(Var) !Term !Term
  | Set _UPK(Id) _UPK(Int)
  deriving (Show, Generic)

data UTerm
  = BApp Term Term
  | BLam Var Term Term
  | BPi Var Term Term
  | BSet Int
  deriving (Show, Generic)

instance Hashable Term
instance Hashable (Description Term)
instance Eq (Description Term)

termCache :: Cache Term
termCache = mkCache
{-# NOINLINE termCache #-}

instance Interned Term where
  type Uninterned Term = UTerm
  data Description Term
    = DApp Id Id
    | DLam Var Id Id
    | DPi Var Id Id
    | DSet Int
    deriving (Show, Generic)

  describe (BApp f a)   = DApp (identity f) (identity a)
  describe (BLam v t e) = DLam v (identity t) (identity e)
  describe (BPi v t e)  = DPi v (identity t) (identity e)
  describe (BSet n)     = DSet n
  
  identify i = go where
      go (BApp f a)   = App i f a
      go (BLam v t e) = Lam i v t e
      go (BPi  v t e) = Pi  i v t e
      go (BSet e)     = Set i e

  cache = termCache

instance Uninternable Term where
  unintern (App _ f a)   = BApp f a
  unintern (Lam _ v t e) = BLam v t e
  unintern (Pi _ v t e)  = BPi v t e
  unintern (Set _ e)     = BSet e

identity :: Term -> Id
identity (App i _ _)   = i
identity (Lam i _ _ _) = i
identity (Pi  i _ _ _) = i
identity (Set i _)     = i

instance Eq Term where
  (==) = (==) `on` identity

instance Ord Term where
  compare = compare `on` identity

app :: Term -> Term -> Term
app a b = intern (BApp a b)

lam :: Var -> Term -> Term -> Term
lam v t e = intern (BLam v t e)

pi :: Var -> Term -> Term -> Term
pi v t e = intern (BPi v t e)

set :: Int -> Term
set i = intern (BSet i)