A small typed markup language based on Haskell syntax

markup.hs

{-# LANGUAGE FlexibleContexts #-}

module Markup where

import Text.Parsec
import qualified Text.Parsec.Token as P
import qualified Text.Parsec.Language as L
import qualified Text.Parsec.Expr as E
import Data.Char
import Data.List (groupBy, nub, (\\))
import Data.Function (on)
import Control.Arrow ((***))
import Control.Applicative ((<$>), (<*>))
import Control.Monad (guard, zipWithM, forM)
import qualified Control.Monad.State as S

data Tm 
  = TmName String
  | TmAtom String 
  | TmStr String 
  | TmInt Integer 
  | TmApp Tm Tm deriving (Show, Eq)

prettyTm :: Tm -> String
prettyTm (TmName s) = s
prettyTm (TmAtom s) = s
prettyTm (TmStr s) = show s
prettyTm (TmInt n) = show n
prettyTm (TmApp t1 t2@(TmApp _ _)) = prettyTm t1 ++ " (" ++ prettyTm t2 ++ ")"
prettyTm (TmApp t1 t2) = prettyTm t1 ++ " " ++ prettyTm t2
    
data Ty
  = TyCon String
  | TyApp Ty Ty
  | TyArr Ty Ty
  | TyVar String deriving (Show, Eq)
  
prettyTy :: Ty -> String
prettyTy (TyCon s) = s
prettyTy (TyVar v) = v
prettyTy (TyApp t1 t2@(TyCon _)) = prettyTy t1 ++ " " ++ prettyTy t2
prettyTy (TyApp t1 t2) = prettyTy t1 ++ " (" ++ prettyTy t2 ++ ")"
prettyTy (TyArr t1 t2) = prettyTy t1 ++ " -> (" ++ prettyTy t2 ++ ")"

unify :: Ty -> Ty -> Either String [(String, Ty)]
unify (TyVar s) (TyVar t) = return []
unify (TyVar s) t = return [(s, t)]
unify t (TyVar s) = return [(s, t)]
unify (TyCon ct1) (TyCon ct2) 
  | ct1 == ct2 = return []
unify (TyApp ty1 ty2) (TyApp ty3 ty4) = (++) <$> unify ty1 ty3 <*> unify ty2 ty4
unify (TyArr ty1 ty2) (TyArr ty3 ty4) = (++) <$> unify ty1 ty3 <*> unify ty2 ty4
unify t1 t2 = Left $ "Cannot unify " ++ prettyTy t1 ++ " with " ++ prettyTy t2
  
subst :: [(String, Ty)] -> Ty -> Ty
subst vs v@(TyVar s) = maybe v id $ lookup s vs
subst vs (TyApp t1 t2) = TyApp (subst vs t1) (subst vs t2)
subst vs (TyArr t1 t2) = TyArr (subst vs t1) (subst vs t2)
subst _ t = t
  
isFunction :: (Eq dom, Eq cod) => [(dom, cod)] -> Bool
isFunction = all ((== 1) . length . nub . map snd) . groupBy ((==) `on` fst)
  
guardS :: (Monad m) => String -> Bool -> m ()
guardS msg b = if b then return () else fail msg
  
specialize :: Ty -> Ty -> Either String Ty
specialize (TyArr t1 t2) t3 = do
  vs <- unify t1 t3
  guardS ("Cannot unify " ++ prettyTy t1  ++ " with " ++ prettyTy t3) $ isFunction vs
  return $ subst vs t2
specialize t _ = Left $ "Expected a function, found " ++ prettyTy t
  
typeOf :: [(String, Ty)] -> Tm -> Either String Ty
typeOf ctx (TmName s) = maybe (Left $ "Unknown name " ++ s) Right $ lookup s ctx
typeOf ctx (TmAtom s) = maybe (Left $ "Unknown constructor " ++ s) Right $ lookup s ctx
typeOf _ (TmInt _) = return (TyCon "Int")
typeOf _ (TmStr _) = return (TyCon "String")
typeOf ctx (TmApp t1 t2) = do
  tyArg <- typeOf ctx t2
  tyFun <- typeOf ctx t1
  flip S.evalState [] $ do 
    tyFun' <- renameAll True tyFun
    tyArg' <- renameAll False tyArg
    return $ specialize tyFun' tyArg'

renameAll :: (Eq a) => a -> Ty -> S.State [((String, a), String)] Ty
renameAll _ t@(TyCon _) = return t
renameAll key (TyVar v) = do
  m <- S.get
  case lookup (v, key) m of
    Nothing -> do
	  let name = unusedName (map snd m)
	  S.modify $ (:) ((v, key), name)
	  return $ TyVar name
    Just name -> return $ TyVar name
renameAll key (TyApp t1 t2) = TyApp <$> renameAll key t1 <*> renameAll key t2
renameAll key (TyArr t1 t2) = TyArr <$> renameAll key t1 <*> renameAll key t2

unusedName :: [String] -> String
unusedName used = head $ map (("t" ++) . show) [0..] \\ used
  
typesOf :: [(String, Ty)] -> [(String, Tm)] -> Either String [(String, Ty)]
typesOf _ [] = return []
typesOf ctx ((name, tm):tms) = do
  ty <- typeOf  ctx tm
  let first = (name, ty)
  rest <- typesOf (first:ctx) tms
  return (first:rest)
  
type Unknown = Int
  
data Kind
  = KindUnknown Unknown
  | KindStar
  | KindArr Kind Kind deriving (Show, Eq)
  
prettyKind :: Kind -> String
prettyKind (KindUnknown n) = "k" ++ show n
prettyKind KindStar = "*"
prettyKind (KindArr k1@(KindArr _ _) k2) = "(" ++ prettyKind k1 ++ ") -> " ++ prettyKind k2
prettyKind (KindArr k1 k2) = prettyKind k1 ++ " -> " ++ prettyKind k2
  
type KindConstraint = (Unknown, Kind)

type SolutionSet = Unknown -> Kind
  
kindCheck :: [Ty] -> Either String [(String, Kind)]
kindCheck = uncurry kindCheck' 
  . (solve *** id) 
  . (concat *** fst) 
  . flip S.runState ([], 0) 
  . flip forM generateConstraints 
  . concatMap applications
  where
  kindCheck' :: Either String SolutionSet -> [(String, Unknown)] -> Either String [(String, Kind)]
  kindCheck' ss unks = do
    ss' <- ss
    return $ map (id *** ss') unks
  
replace :: Unknown -> Kind -> Kind -> Kind
replace i x u@(KindUnknown j) | i == j = x
                              | otherwise = u
replace _ _ KindStar = KindStar
replace i x (KindArr k1 k2) = KindArr (replace i x k1) (replace i x k2)

occursCheck :: Unknown -> Kind -> Either String ()
occursCheck i (KindUnknown j) | i == j = return ()
occursCheck i x = occursCheck' i x where
  occursCheck' i (KindUnknown j) | i == j = fail "Occurs check failed during kind checking"
  occursCheck' i (KindArr k1 k2) = occursCheck' i k1 >> occursCheck' i k2
  occursCheck' _ _ = return ()

unifyKinds :: Kind -> Kind -> Either String [KindConstraint]
unifyKinds (KindUnknown i) x = return [(i, x)]
unifyKinds x (KindUnknown i) = return [(i, x)]
unifyKinds KindStar KindStar = return  []
unifyKinds (KindArr k1 k2) (KindArr k3 k4) = (++) <$> unifyKinds k1 k3 <*> unifyKinds k2 k4
unifyKinds k1 k2 = fail $ "Cannot unify " ++ show k1 ++ " with " ++ show k2
  
solve :: [KindConstraint] -> Either String SolutionSet
solve cs = solve' cs KindUnknown where
  solve' [] ss = return ss
  solve' (c@(i, x):cs) ss = do
    occursCheck i x
    cs' <- substConstraints c cs
    let ss' = substSolutionSet c ss
    solve' cs' ss'
  substConstraints c [] = return []
  substConstraints (i, x) ((j, y): cs) 
    | i == j = (++) <$> substConstraints (i, x) cs <*> unifyKinds x y
    | otherwise = (:) (j, replace i x y) <$> substConstraints (i, x) cs
  substSolutionSet (i, x) = (.) (replace i x)
  
applications :: Ty -> [Ty]
applications t = applications' t [] where
  applications' (TyArr t1 t2) = applications' t1 . applications' t2
  applications' t = (:) t
  
generateConstraints :: (Monad m, S.MonadState ([(String, Unknown)], Unknown) m) => Ty -> m [KindConstraint]
generateConstraints ty = do
  (cs, n) <- generateConstraints' ty
  return $ (n, KindStar) : cs
  where
  generateConstraints' :: (Monad m, S.MonadState ([(String, Unknown)], Unknown) m) => Ty -> m ([KindConstraint], Unknown)
  generateConstraints' (TyVar v) = do
    (m, _) <- S.get
    case lookup v m of
      Nothing -> do
        n <- newUnknown
        S.modify $ (:) (v, n) *** id
        return ([], n)
      Just unk -> return ([], unk)
  generateConstraints' (TyCon s) = do
    (m, _) <- S.get
    case lookup s m of
      Nothing -> do
        n <- newUnknown
        S.modify $ (:) (s, n) *** id
        return ([], n)
      Just unk -> return ([], unk)
  generateConstraints' (TyApp ty1 ty2) = do
    (c1, n1) <- generateConstraints' ty1
    (c2, n2) <- generateConstraints' ty2
    thisName <- newUnknown
    let newConstraint = (n1, KindArr (KindUnknown n2) (KindUnknown thisName))
    return $ ((newConstraint:c1) ++ c2, thisName)
    
newUnknown :: (Monad m, S.MonadState (a, Unknown) m) => m Unknown
newUnknown = S.modify (id *** (+1)) >> S.get >>= return . snd
     
languageDef :: P.LanguageDef st
languageDef = L.haskellStyle 
  { P.opStart         = fail "Operators are not supported"
  , P.opLetter        = fail "Operators are not supported"
  , P.reservedNames   = []
  , P.reservedOpNames = []
  , P.caseSensitive   = True }
  
lexer             = P.makeTokenParser languageDef 
identifier        = P.identifier lexer
whiteSpace        = P.whiteSpace lexer
stringLiteral     = P.stringLiteral lexer 
integer           = P.integer lexer 
symbol            = P.symbol lexer 
lexeme            = P.lexeme lexer 
parens            = P.parens lexer 
semi              = P.semi lexer 

identu :: Parsec String u String
identu = lookAhead upper >> identifier
  
identl :: Parsec String u String
identl = lookAhead lower >> identifier
  
tmName :: Parsec String u Tm
tmName = lexeme $ identl >>= return . TmName
  
tmAtom :: Parsec String u Tm
tmAtom = lexeme $ identu >>= return . TmAtom

tmStr :: Parsec String u Tm
tmStr = lexeme $ stringLiteral >>= return . TmStr
  
tmInt :: Parsec String u Tm
tmInt = lexeme $ integer >>= return . TmInt

tmBracket :: Parsec String u Tm
tmBracket = lexeme $ parens tm
  
tm :: Parsec String u Tm
tm = E.buildExpressionParser 
     [ [ E.Infix (return TmApp) E.AssocLeft ]
	 , [ E.Infix (((.) TmApp . TmApp) <$> between (symbol "`") (symbol "`") tmAtom) E.AssocLeft ]
     , [ E.Infix (symbol "$" >> return TmApp) E.AssocRight ]  ]
     (tmBracket <|> tmStr <|> tmInt <|> try tmName <|> tmAtom)
	
assignment :: Parsec String u (String, Tm)
assignment = do
  name <- identl
  symbol "="
  t <- tm
  return (name, t)

doc :: Parsec String u [(String, Tm)]
doc = do
  whiteSpace 
  ds <- sepBy assignment (lexeme semi)
  eof
  return ds
	 
tyVar :: Parsec String u Ty
tyVar = lexeme $ identl >>= return . TyVar

tyCon :: Parsec String u Ty
tyCon = lexeme $ identu >>= return . TyCon
  
tyBracket :: Parsec String u Ty
tyBracket = parens ty
  
ty :: Parsec String u Ty
ty = E.buildExpressionParser 
     [ [ E.Infix (return TyApp) E.AssocLeft ]
     , [ E.Infix (symbol "->" >> return TyArr) E.AssocRight ] ]
     (tyBracket <|> try tyVar <|> tyCon)
	
decl :: Parsec String u (String, Ty)
decl = do
  name <- identu
  symbol "::"
  t <- ty
  return (name, t)

schema :: Parsec String u [(String, Ty)]
schema = do
  whiteSpace 
  ds <- sepBy decl (lexeme semi)
  eof
  return ds

parseErr p = either (Left . show) Right . parse p ""

exampleSchema = unlines 
  [ "Nil :: List Z a;"
  , "Singleton :: a -> List (S Z) a;"
  , "Cons :: a -> List n a -> List (S n) a;"
  , "Reverse :: List n a -> List n a;"
  , "Zip :: List n a -> List n b -> List n (Pair a b);"
  , "Tip :: Tree Z a;"
  , "Branch :: Tree n a -> a -> Tree n a -> Tree (S n) a" ]
  
example1 = "x = Zip (Cons 1 $ Cons 2 Nil) (Cons 1 Nil)" 
example2 = "x = Zip (Cons 1 $ Reverse $ Cons 2 $ Cons 3 Nil) (Cons \"a\" $ Cons \"b\" $ Singleton \"c\")" 
example3 = "x = Branch (Branch Tip 1 Tip) 2 (Branch Tip 3 Tip)" 
example4 = "x = Branch (Branch Tip 1 Tip) 2 Tip" 
example5 = "x = Cons (Branch Tip 1 Tip) Nil" 
example6 = "x = 1 `Cons` (2 `Cons` (3 `Cons` Nil))" 
example7 = "x = 1 `Cons` Nil `Zip` (2 `Cons` Nil)" 
example8 = "x = Nil `Cons` Nil `Zip` ((Cons 1 $ Nil) `Cons` Nil)" 
example9 = unlines 
  [ "x = Nil;"
  , "y = Cons x Nil;"
  , "z = Cons y Nil" ]
  
test sch val = do
  s <- parseErr schema sch
  ks <- kindCheck $ map snd s
  d <- parseErr doc val
  tys <- typesOf s d
  return $ map (id *** prettyKind) ks

relationSchema = unlines 
  [ "UserID         :: Column User PK;"
  , "UserName       :: Column User String;"
  , "UserCompanyID  :: Column User PK;"
  , "CompanyID      :: Column Company PK;"
  , "CompanyName    :: Column Company String;"
  , "Table          :: HList (Column t) -> Query (Table t);"
  , "Join           :: Query r1 -> Query r2 -> Expr (Cross r1 r2) Bool -> Query (Cross r1 r2);"
  , "Column         :: Column t ty -> Expr (Table t) ty;"
  , "Left           :: Expr t1 ty -> Expr (Cross t1 t2) ty;"
  , "Right          :: Expr t2 ty -> Expr (Cross t1 t2) ty;"
  , "Eq             :: Expr t ty -> Expr t ty -> Expr t Bool;"
  , "Nil            :: HList f;"
  , "Cons           :: f t -> HList f -> HList f" ]

relationExample = unlines 
  [ "query = Join"
  , "          (Table (Cons UserName $ Cons UserCompanyID Nil))"
  , "          (Table (Cons CompanyID $ Cons CompanyName Nil))"
  , "          ((Left $ Column UserCompanyID) `Eq` (Right $ Column CompanyID))" ]