Enumerates possible types of untyped terms in system F

sysF.hs

{-# LANGUAGE LambdaCase, BlockArguments #-}
module Main where
import Data.Maybe
import Data.Bifunctor
import Control.Applicative
import Control.Monad
-- weighted-search package
import qualified Control.Monad.WeightedSearch as W

greeks :: [Char]
greeks = "αβγδεζηθικλμνξοπρστυφχψω"

latins :: [Char]
latins = "xyzwuvabcdefghijklmnopqrst"

type Var = String

data Type = TVar Int {- dbL -} | Arr Type Type | ForAll Type deriving (Eq)

showType :: Int -> Int -> Type -> ShowS
showType k p = \case
  -- Special shorthands
  ty | ty == falseType k -> showString "⊥"
     | ty == trueType k -> showString "⊤"
     | ty == boolType k -> showString "Bool"
     | ty == natType k -> showString "Nat"
  TVar j -> showChar (greeks !! j)
  Arr dom cod -> showParen (p > 0) $
    showType k 1 dom . showString " → " . showType k 0 cod
  ForAll ty -> showParen (p > 0) $
    showString "∀" . showChar (greeks !! k) . showString ". " .
    showType (k+1) 0 ty

instance Show Type where
  showsPrec = showType 0

falseType :: Int -> Type
falseType k = ForAll (TVar k)

trueType :: Int -> Type
trueType k = ForAll (Arr (TVar k) (TVar k))

boolType :: Int -> Type
boolType k = ForAll (Arr (TVar k) (Arr (TVar k) (TVar k)))

natType :: Int -> Type
natType k = ForAll (Arr (Arr (TVar k) (TVar k)) (Arr (TVar k) (TVar k)))

weaken :: Int -> Type -> Type
weaken k = \case
  TVar j -> if j < k then TVar j else TVar (j + 1)
  Arr t1 t2 -> Arr (weaken k t1) (weaken k t2)
  ForAll t -> ForAll (weaken k t)

-- | subst k r substitutes the k-th variable for sb in a type with r variables
subst :: Int -> Int -> Type -> Type -> Type
subst k r sb = \case
  TVar j
    | j < k -> TVar j
    | j == k -> sb
    | otherwise -> TVar (j - 1)
  Arr t1 t2 -> Arr (subst k r sb t1) (subst k r sb t2)
  ForAll t -> ForAll (subst k (r + 1) (weaken r sb) t)

data Term
  = Var Var
  | Lam Var Type Term | App Term Term
  | TLam Term | TApp Term Type

showTerm :: Int -> Int -> Term -> ShowS
showTerm k p = \case  -- TODO show church numerals
  Var x -> showString x
  Lam x ty tm -> showParen (p > 0) $
    showString "λ" .
    showParen True (showString x . showString " : " . showType k 0 ty) .
    showString ". " .
    showTerm k 0 tm
  App f t -> showParen (p > 1) $
    showTerm k 1 f . showString " " . showTerm k 2 t
  TLam tm -> showParen (p > 0) $
    showString "Λ" . showChar (greeks !! k) . showString ". " .
    showTerm (k+1) 0 tm
  TApp f t -> showParen (p > 1) $
    showTerm k 1 f . showString " [" . showType k 0 t . showString "]"

instance Show Term where
  showsPrec = showTerm 0

weaken' :: Int -> Term -> Term
weaken' k = \case
  Var x -> Var x
  Lam x ty tm -> Lam x (weaken k ty) (weaken' k tm)
  App t1 t2 -> App (weaken' k t1) (weaken' k t2)
  TLam t -> TLam (weaken' k t)
  TApp tm ty -> TApp (weaken' k tm) (weaken k ty)

fresh :: [Var] -> Var
fresh vars = head $ filter (`notElem` vars) vs
  where
    vs = [c : n | n <- "" : map show [0::Int .. 9], c <- latins]

data Raw
  = RVar Var
  | RLam Var Raw | RApp Raw Raw

-- | Generate all types under a given type context
allTypes :: Int -> W.T Integer Type
allTypes k = asum [return (TVar i) | i <- [0..k-1]]
  <|> W.weight 2 (ForAll <$> allTypes (k+1))
  <|> W.weight 1 (Arr <$> allTypes k <*> allTypes k)

-- The "natural type" of things, inserting nothing on the outside
infer :: Int -> [(Var, Type)] -> Raw -> W.T Integer (Term, Type)
infer tv ctx = \case
  RVar x -> let Just xty = lookup x ctx in return (Var x, xty)
  RLam x tm -> do
    argty <- allTypes tv
    (tm', codty) <- infer tv ((x,argty):ctx) tm
    return (Lam x argty tm', Arr argty codty)
  RApp f t -> do
    (f', fty) <- infer tv ctx f
    -- check all instantiations
    (args, fty') <- instantiate tv fty
    case fty' of
      Arr farg fcod -> do
        t' <- check tv ctx t farg
        return (App (foldl TApp f' args) t', fcod)
      _ -> empty

instantiate :: Int -> Type -> W.T Integer ([Type], Type)
instantiate k = \case
  ForAll ty -> do
    arg <- allTypes k
    let ty' = subst k (k + 1) arg ty
    (args, ty'') <- instantiate k ty'
    return (arg : args, ty'')
  ty -> return ([], ty)

check :: Int -> [(Var, Type)] -> Raw -> Type -> W.T Integer Term
check tv ctx tm (ForAll ty) =
  -- insert as many type lambdas as we want
  TLam <$> check (tv + 1) (map (second (weaken tv)) ctx) tm ty
check tv ctx tm ty = case tm of
  RLam x tm -> case ty of
    Arr dom cod -> Lam x dom <$> check tv ((x,dom):ctx) tm cod
    _ -> empty -- no more ForAll to deal with
  _ -> do
    -- infer the natural type, and figure out what type applications we do
    (tm', ty') <- infer tv ctx tm
    let (k, rty') = rip ty'
    sub <- matchType tv (tv + k) rty' ty
    return $ foldl' TApp tm' $ map (fromMaybe (falseType tv)) sub

rip :: Type -> (Int, Type)
rip (ForAll ty) = first (+1) (rip ty)
rip ty = (0, ty)

-- | Assign variables in [tvmin, tvmax) so that the first type matches the second
matchType :: Int -> Int -> Type -> Type -> W.T Integer [Maybe Type]
matchType tvmin tvmax (TVar i) goal
  | tvmin <= i && i < tvmax = return [
    if j == i then Just goal else Nothing |
    j <- [tvmin..tvmax - 1]
  ]
  | goal == TVar (i - tvmax + tvmin) = return [
    Nothing |
    _ <- [tvmin..tvmax - 1]
  ]
  | otherwise = empty
matchType tvmin tvmax (Arr t1 t2) (Arr g1 g2) = do
  s1 <- matchType tvmin tvmax t1 g1
  s2 <- matchType tvmin tvmax t2 g2
  reconcile s1 s2
  where
    reconcile [] [] = return []
    reconcile (Just ty1 : s1) (Nothing : s2) = (Just ty1 :) <$> reconcile s1 s2
    reconcile (Nothing : s1) (Just ty2 : s2) = (Just ty2 :) <$> reconcile s1 s2
    reconcile (Just ty1 : s1) (Just ty2 : s2)
      | ty1 == ty2 = (Just ty1 :) <$> reconcile s1 s2
    reconcile _ _ = empty
matchType tvmin tvmax (ForAll t) (ForAll g) = matchType tvmin tvmax t g
matchType _ _ _ _ = empty

-- | Generate normal forms of a given type
genNormal, genNeutral, genVar :: Int -> [(Var, Type)] -> Type -> W.T Integer Term

genNormal tv ctx = \case
  TVar j -> genNeutral tv ctx (TVar j)
  Arr dom cod -> do
    let x = fresh (map fst ctx)
    Lam x dom <$> genNormal tv ((x,dom):ctx) cod
  ForAll ty -> do
    TLam <$> genNormal (tv + 1) (map (second (weaken tv)) ctx) ty

genNeutral tv ctx ty = genVar tv ctx ty
  <|> W.weight 2 do
    dom <- allTypes tv
    arg <- genNormal tv ctx dom
    (`App` arg) <$> genNeutral tv ctx (Arr dom ty)
  <|> W.weight 3 do
    arg <- allTypes tv
    newTy <- extract tv tv arg ty
    (`TApp` arg) <$> genNeutral tv ctx (ForAll newTy)

genVar _ ctx ty = asum do
  (x, ty') <- ctx
  guard (ty == ty')
  return (return (Var x))

-- extract k r sub ty replaces occurrences of sub in ty by the k-th variable
-- where ty has r variables (and will have (r+1) variables)
extract :: Int -> Int -> Type -> Type -> W.T Integer Type
extract k r sub ty
  | sub == ty = return (TVar k) <|> W.weight 2 (extract' k r sub ty)
  | otherwise = extract' k r sub ty
extract' k r sub ty = case ty of
  TVar j
    | j < k -> return $ TVar j
    | otherwise -> return $ TVar (j + 1)
  Arr t1 t2 -> Arr <$> extract k r sub t1 <*> extract k r sub t2
  ForAll ty -> ForAll <$> extract k (r+1) (weaken r sub) ty

program = infer 0 [] $
  RLam "x" (RApp (RVar "x") (RVar "x"))

-- program = check 0 []
--   (RLam "x" (RApp (RVar "x") (RVar "x")))
--   (Arr (falseType 0) (falseType 0))

-- program = genNormal 0 [] (Arr (boolType 0) (boolType 0))

main :: IO ()
main = do
  forM_ (take 20 (W.toList program)) \(tm, ty) -> do
    print tm
    print ty
    putStrLn ""