system wau preview

stack.yaml

resolver: lts-16.6
packages:
- .

extra-deps:
- unsafeperformst-0.9.2@sha256:de35debea25dd6a500532a57efe82c625e94dd93b8da323a2b93ddb288c7673b,734

system-wau.cabal

cabal-version: 1.12

-- This file has been generated from package.yaml by hpack version 0.33.1.
--
-- see: https://github.com/sol/hpack
--
-- hash: 1fc780c52d297504d7e8bf6a820619adf3268f07676703e6cc8eb53dbcc73b26

name:           system-wau
version:        0.1.0.0
description:    Please see the README on GitHub at <https://github.com/johnchandlerburnham/system-wau#readme>
homepage:       https://github.com/johnchandlerburnham/system-wau#readme
bug-reports:    https://github.com/johnchandlerburnham/system-wau/issues
author:         John C. Burnham
maintainer:     jcb@johnchandlerburnham.com
copyright:      2019 John Burnham
license:        BSD3
license-file:   LICENSE
build-type:     Simple
extra-source-files:
    README.md
    ChangeLog.md

source-repository head
  type: git
  location: https://github.com/johnchandlerburnham/system-wau

library
  exposed-modules:
      Wau
  other-modules:
      Lib
      Parse
      Print
      Paths_system_wau
  hs-source-dirs:
      src
  default-extensions: OverloadedStrings MultiWayIf
  build-depends:
      base >=4.7 && <5
    , containers
    , megaparsec
    , mtl
    , text
    , unsafeperformst
  default-language: Haskell2010

Wau.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE MultiWayIf #-}
module Wau where

import           Data.Map                       (Map)
import qualified Data.Map                       as M
import           Data.Maybe                     (isJust)
import           Data.Set                       (Set)
import qualified Data.Set                       as Set
import           Data.STRef
import           Data.Text                      (Text)
import qualified Data.Text                      as T
import qualified Data.Text.IO                   as TIO
import           Data.Void
import           Debug.Trace

import           Control.Monad.Identity
import           Control.Monad.ST
import           Control.Monad.ST.UnsafePerform
import           Control.Monad.Except

import           System.Exit

import           Text.Megaparsec                hiding (State)
import           Text.Megaparsec.Char
import qualified Text.Megaparsec.Char.Lexer     as L

type Name = Text
data Loc = Loc { _from :: Int, _upto :: Int } deriving Show

noLoc = Loc 0 0

data Term
  = Var Loc Name Int
  | Lam Loc Bool Name (Term -> Term)
  | App Loc Bool Term Term
  | All Loc Bool Name Name Term (Term -> Term -> Term)
  | Fix Loc Name (Term -> Term)
  | Let Loc Name Term (Term -> Term)
  | Ref Loc Name
  | Typ Loc
  | Ann Loc Bool Term Term

data Expr = Expr { _name :: Name, _type :: Term, _term :: Term }
data Defs = Defs { _defs :: (Map Name Expr)}

emptyDefs = Defs M.empty

type Ctx = [(Text,Term)]

find :: Ctx -> ((Text,Term) -> Bool) -> Maybe ((Text,Term),Int)
find ctx f = go ctx 0
  where
    go [] _     = Nothing
    go (x:xs) i = if f x then Just (x,i) else go xs (i+1)

compText :: Term -> Int -> Text
compText term dep = let go = compText in case term of
  Var _ n x  -> if x < 0
    then T.concat ["^", T.pack $ show $ dep + x]
    else T.concat ["#", T.pack $ show x]
  Typ _      -> "*"
  Ref _ n          -> T.concat ["&", n] 
  All _ e _ _ h b ->
    let all  = if e then "∀0" else "∀ω"
        bind = go h dep
        s    = Var noLoc "" (0-dep-1)
        x    = Var noLoc "" (0-dep-2)
        body = go (b s x) (dep + 2)
     in T.concat [all,bind,body]
  Fix _ _ b ->
    let body = go (b (Var noLoc "" (0-dep-1))) (dep+1)
     in T.concat ["μ", body]
  Lam _ e n b          ->
    let lam = if e then "λ0" else "λω"
        body = go (b (Var noLoc "" (0-dep-1))) (dep+1)
    in T.concat [lam, body]
  App _ e f a      -> 
    let app = if e then "@0" else "@ω"
     in T.concat [app, go f dep, go a dep]
  Let _ _ x b        ->
    let expr = go x dep
        body = go (b (Var noLoc "" (0-dep-1))) (dep+1)
     in T.concat ["$",expr,body]
  Ann _ d x t          -> go x dep

hash :: Term -> Text
hash t = compText t 0

printTerm :: Term -> Text
printTerm t = let go = printTerm in case t of
  Var _ "" i       -> T.concat ["^", T.pack $ show i]
  Var _ n i        -> n
  Typ _            -> "*"
  Ref _ n          -> T.concat ["#",n]
  All _ e s n h b  ->
    let eras = if e then "0 " else ""
        body = go (b (Var noLoc s 0) (Var noLoc n 0))
     in T.concat ["∀", s, "(", eras, n, ":", go h, ") ", body]
  Fix _ n b ->
    let body = go (b (Var noLoc n 0))
     in T.concat ["μ",n,".",body]
  Lam _ e n b      ->
    let eras = if e then "0 " else ""
        body = go (b (Var noLoc n 0))
    in T.concat ["λ", eras, n, ". ", body]
  App _ True  f a      -> T.concat ["(0 ", go f, " ", go a, ")"]
  App _ False f a      -> T.concat ["(", go f, " ", go a, ")"]
  Let _ n x b        -> let body = go (b (Var noLoc n 0)) in 
    T.concat ["$", n, "=", go x, ";", body]
  Ann _ d x t          -> T.concat ["(", go x, " :: ", go t, ")"]

printExpr :: Expr -> Text
printExpr (Expr n t x) = T.concat [n,": ",printTerm t,"\n  ",printTerm x]

printDefs :: Defs -> Text
printDefs (Defs ds) = go $ snd <$> M.toList ds
  where 
    go [] = ""
    go (x:[]) = printExpr x
    go (x:xs) = T.concat [printExpr x, "\n", go xs]

instance Show Term where
  show = T.unpack . printTerm

instance Show Expr where
  show = T.unpack . printExpr

instance Show Defs where
  show = T.unpack . printDefs

-- Parser

type Parser = ParsecT Void Text Identity

name :: Bool -> Parser Text
name empty = (if empty then takeWhileP else takeWhile1P)
  (Just "a name (alphanumeric,'_','.')") (flip elem $ nameChar)
  where
    nameChar = ['0'..'9'] ++ ['a'..'z'] ++ ['A'..'Z'] ++ "_"

spaceC :: Parser ()
spaceC = L.space space1 (L.skipLineComment "//") (L.skipBlockComment "#" "#")

symbol :: Text -> Parser Text
symbol = L.symbol spaceC

-- The term parser
-- TODO: proper error printing instead of the `\n` hack to
-- make the ShowErrorComponent print nicely
-- (Perhaps by using a typerep proxy plus a new typeclass instance)
term :: Defs -> [Name] -> Parser (Ctx -> Term)
term ds ns = do
  from <- getOffset
  t    <- choice $
    [ label "\n - the type of types: \"*\"" $
       symbol "*" >> (return $ \ctx -> Typ (Loc from (from+1)))
    , label "\n - a forall: \"∀self(0 x: A) B\", \"∀self(x: A) B\"" $ do
        string "∀"
        self <- name True <* symbol "("
        eras <- isJust <$> (optional $ symbol "0 ")
        name <- name True <* spaceC <* symbol ":"
        bind <- term ds (self:ns) <* symbol ")"
        symbol "->"
        body <- term ds (name:self:ns)
        upto <- getOffset
        let l = Loc from upto
        return $ \ctx ->
          All l eras self name (bind ctx) (\s x -> body ((name,x):(self,s):ctx))
    , label "\n - a lambda: \"λ0 x. b\", \"λx. b\"" $ do
        from <- getOffset
        string "λ"
        eras <- isJust <$> (optional $ symbol "0 ")
        name <- name True <* spaceC <* symbol "."
        body <- term ds (name:ns)
        upto <- getOffset
        let l = Loc from upto
        return $ \ctx -> Lam l eras name (\x -> body ((name,x):ctx))
    , label "\n - an application: \"(0 f a)\", \"(f a)\"" $ do
        string "("
        eras <- isJust <$> (optional $ symbol "0 ")
        func <- term ds ns
        argm <- term ds ns
        symbol ")"
        upto <- getOffset
        let l = Loc from upto
        return $ \ctx -> App l eras (func ctx) (argm ctx)
    , label "\n - a fixpoint: \"μx. (f x)\"" $ do
        string "μ"
        name <- name False <* spaceC <* symbol "."
        body <- term ds (name:ns)
        upto <- getOffset
        let l = Loc from upto
        return $ \ctx -> Fix l name (\x -> body ((name,x):ctx))
    , label "\n - a definition: \"$x = y; b\"" $ do
        string "$"
        name <- name True <* spaceC <* symbol "="
        expr <- term ds ns <* symbol ";"
        body <- term ds (name:ns)
        upto <- getOffset
        let l = Loc from upto
        return $ \ctx -> Let l name (expr ctx) (\x -> body ((name,x):ctx))
    , label "\n - a type annotation: \"{x:A}\"" $ do
        symbol "{"
        typ_ <- term ds ns
        symbol ":"
        expr <- term ds ns
        upto <- getOffset
        symbol "}"
        let l = Loc from upto
        return $ \ctx -> Ann l False (expr ctx) (typ_ ctx)
    , label "\n - a reference, either global or local: \"x\"" $ do
        name <- name False
        upto <- getOffset
        let l = Loc from upto
        case find ns (\(n,i) -> n == name) of
          Just ((n,t),i) -> return $ \ctx -> t
          Nothing -> case deref name defs of
            Nothing -> fail $ "Undefined Reference: " ++ T.unpack name
            Just h  -> return $ \ctx -> Ref l name
    ]
  spaceC
  return t

expr :: Defs -> Parser Expr
expr defs = do
  name <- name False <* spaceC <* symbol ":"
  typ_ <- ($ []) <$> term defs []
  term <- ($ []) <$> term defs []
  return $ Expr name typ_ term

defs :: Parser Defs
defs = Defs . M.fromList . fmap (\d -> (_name d, d)) <$> exprs emptyDefs
  where
    exprs :: Defs -> Parser [Expr]
    exprs ds = next ds <|> (spaceC >> return [])

    next :: Defs -> Parser [Expr]
    next ds = do
      spaceC
      x <- expr ds
      (x:) <$> (exprs $ Defs $ M.insert (_name x) x (_defs ds))


deref :: Name -> Defs -> Maybe Expr
deref n m = (_defs m) M.!? n

parseTerm :: Text -> IO Term
parseTerm txt = case parse (term emptyDefs []) "" txt of
  Left  e -> putStr (errorBundlePretty e) >> exitFailure
  Right t -> return (t $ [])

parseFile :: FilePath -> IO Defs
parseFile file = do
  txt <- TIO.readFile file
  case parse defs file txt of
    Left  e -> putStr (errorBundlePretty e) >> exitFailure
    Right m -> return m

parseExpr :: FilePath -> Name -> IO Expr
parseExpr fileName termName = do
  defs <- parseFile fileName
  case deref termName defs of
    Just x  -> return x
    Nothing -> putStr "Undefined Reference" >> exitFailure

normExpr :: FilePath -> Name -> IO Term
normExpr fileName termName = do
  defs <- parseFile fileName
  case deref termName defs of
    Just x  -> return $ normalize (_term x) defs True
    Nothing -> putStr "Undefined Reference" >> exitFailure


-- Evaluation

unroll :: Term -> Term
unroll t = case t of
  Fix loc name body -> body (Fix loc name body)
  _ -> t

reduce :: Term -> Defs -> Bool -> Term
reduce term (Defs defs) erase = go term
  where
    go term = case term of
      Var l n d        -> Var l n d
      Ref l n          -> case defs M.!? n of
        Just (Expr _ _ got)             -> go got
        Nothing                         -> term
      All _ _ _ _ _ _  -> term
      Typ _            -> Typ noLoc
      Fix l n b        -> go $ unroll term
      Lam _ e n b      ->
          if e && erase then go (b (Lam noLoc False "" (\x -> x))) else term
      App _ e f a      -> if e && erase then go f else case go f of
        Lam _ e n b -> go (b a)
        x           -> term
      Let _ n x b      -> go (b x)
      Ann _ _ x t      -> go x

-- Normalize
normalize :: Term -> Defs -> Bool -> Term
normalize term defs erased = runST (top term)
  where
    top :: Term -> ST s Term
    top term = do
      seen <- newSTRef (Set.empty)
      go term seen

    go :: Term -> (STRef s (Set Text)) -> ST s Term
    go term seen = {- trace "go" $ -} do
      let norm  = reduce term defs erased
      let termH = hash term
      let normH = hash norm
      -- traceM $ concat ["term: ",show term, " hash: ",show termH]
      -- traceM $ concat ["norm: ",show norm, " hash: ",show normH]
      seen' <- readSTRef seen
      -- traceM $ concat ["seen: ",show seen']
      if | (termH `Set.member` seen' || normH `Set.member` seen') -> return norm
         | otherwise -> do
           modifySTRef' seen ((Set.insert termH) . (Set.insert normH))
           case norm of
             Var l n d        -> {- trace "norm Var" $ -} return $ Var l n d
             Ref l n          -> {- trace "norm Ref" $ -} return $ Ref l n
             Typ l            -> {- trace "norm Typ" $ -} return $ Typ l
             All _ e s n h b  -> {- trace "norm All" $ -} do
               bind <- go h seen
               return $ All noLoc e s n bind (\s x -> unsafePerformST $ go (b s x) seen)
             Lam _ e n b      -> {-trace "norm Lam" $ traceShow norm $-} do
               return $ Lam noLoc e n (\x -> unsafePerformST $ go (b x) seen)
             App _ e f a      -> {-trace "norm App" $ -} do
              func <- go f seen
              argm <- go a seen
              return $ App noLoc e func argm
             Let _ n x b      -> {-trace "norm Let" $-} go (b x) seen
             Fix l n b        -> {-trace "norm Let" $-} go norm seen
             Ann _ _ x t      -> {-trace "norm Ann" $-} go x seen

-- Typechecking
equal :: Term -> Term -> Defs -> Int -> Bool
equal a b defs dep = runST $ top a b dep
  where
    top :: Term -> Term -> Int -> ST s Bool
    top a b dep = do
      seen <- newSTRef (Set.empty)
      go a b dep seen

    go :: Term -> Term -> Int -> STRef s (Set (Text,Text)) -> ST s Bool
    go a b dep seen = do
      let var d = Var noLoc "" (fromIntegral d)
      let a1 = reduce a defs False
      let b1 = reduce b defs False
      let ah = hash a1
      let bh = hash b1
      -- traceM $ show a
      -- traceM $ show b
      s' <- readSTRef seen
      if | (ah == bh)              -> return True
         | (ah,bh) `Set.member` s' -> return True
         | (bh,ah) `Set.member` s' -> return True
         | otherwise -> do
             modifySTRef' seen ((Set.insert (ah,bh)) . (Set.insert (bh,ah)))
             case (a1,b1) of
               (All _ ae as _ ah ab, All _ be bs _ bh bb) -> do
                 let a1_body = ab (var dep) (var (dep + 1))
                 let b1_body = bb (var dep) (var (dep + 1))
                 let eras_eq = ae == be
                 let self_eq = as == bs
                 bind_eq <- go ah bh dep seen
                 body_eq <- go a1_body b1_body (dep+2) seen
                 return $ eras_eq && self_eq && bind_eq && body_eq
               (Lam _ ae _ ab, Lam _ be _ bb) -> do
                 let a1_body = ab (var dep)
                 let b1_body = bb (var dep)
                 let eras_eq = ae == be
                 body_eq <- go a1_body b1_body (dep+1) seen
                 return $ eras_eq && body_eq
               (App _ ae af aa, App _ be bf ba) -> do
                 let eras_eq = ae == be
                 func_eq <- go af bf dep seen
                 argm_eq <- go aa ba dep seen
                 return $ eras_eq && func_eq && argm_eq
               (Let _ _ ax ab, Let _ _ bx bb) -> do
                 let a1_body = ab (var dep)
                 let b1_body = bb (var dep)
                 expr_eq <- go ax bx dep seen
                 body_eq <- go a1_body b1_body (dep+1) seen
                 return $ expr_eq && body_eq
               (Ann _ _ ax _, Ann _ _ bx _) -> go ax bx dep seen
               _ -> return False

data CheckErr = CheckErr Loc Ctx Text deriving Show

check :: Term -> Term -> Defs -> Ctx -> Except CheckErr Bool
check trm typ defs ctx = do
  --traceM $ "check"
  --traceM $ show trm
  --traceM $ show typ
  let trm  = unroll trm
  let typv = reduce typ defs False
  let var n l = Var noLoc "" (fromIntegral l)
  case trm of
    Lam trm_loc trm_eras trm_name trm_body -> case typv of
      All _ typ_eras typ_self typ_name typ_bind typ_body ->
        if typ_eras /= trm_eras 
        then throwError (CheckErr trm_loc ctx "Type mismatch")
        else do
          let self_var = Ann noLoc True trm typ
          let name_var = Ann noLoc True (var trm_name (length ctx + 1)) typ_bind
          let body_typ = typ_body self_var name_var
          let body_ctx = (trm_name,typ_bind):ctx
          check (trm_body name_var) body_typ defs body_ctx
      _  -> do
      --traceM $ show typv
        throwError (CheckErr trm_loc ctx "Lamda has non-function type")
    Let trm_loc trm_name trm_expr trm_body -> do
      expr_typ <- infer trm_expr defs ctx
      let expr_var = Ann noLoc True (var trm_name (length ctx + 1)) expr_typ
      let body_ctx = (trm_name,expr_typ):ctx
      check trm_expr typ defs body_ctx
    _ -> do
      infr <- infer trm defs ctx
      let eq   = equal typ infr defs (length ctx)
      if eq
      then return True
      else throwError (CheckErr noLoc ctx "bad")

infer :: Term -> Defs -> Ctx -> Except CheckErr Term
infer trm defs ctx = do
  let var n l = Var noLoc "" (fromIntegral l)
  case trm of
    Var l n d -> return trm
    Ref l n -> case (_defs defs) M.!? n of
      Just (Expr _ t _) -> return t
      Nothing           -> throwError (CheckErr l ctx (T.concat ["Undefined Reference ", n]))
    Typ l   -> return $ Typ l
    App trm_loc trm_eras trm_func trm_argm -> do
      func_typ <- (\x -> reduce x defs False) <$> infer trm_func defs ctx
      case func_typ of
        All ftyp_loc ftyp_eras ftyp_self_ ftyp_name ftyp_bind ftyp_body -> do
          let self_var = Ann noLoc True trm_func func_typ
          let name_var = Ann noLoc True trm_argm ftyp_bind
          check trm_argm ftyp_bind defs ctx
          let trm_typ = ftyp_body self_var name_var
          if trm_eras /= ftyp_eras
          then throwError $ CheckErr trm_loc ctx "Mismatched Erasure"
          else return trm_typ
        _ -> throwError $ CheckErr trm_loc ctx "Non-function application"
    Let trm_loc trm_name trm_expr trm_body -> do
      expr_typ <- infer trm_expr defs ctx
      let expr_var = Ann noLoc True (var trm_name (length ctx + 1)) expr_typ
      let body_ctx = (trm_name,expr_typ):ctx
      infer (trm_body expr_var) defs body_ctx
    All trm_loc trm_eras trm_self trm_name trm_bind trm_body -> do
      let self_var = Ann noLoc True (var trm_self $ length ctx) trm
      let name_var = Ann noLoc True (var trm_name $ length ctx + 1) trm_bind
      let body_ctx = (trm_name,trm_bind):(trm_self,trm):ctx
      check trm_bind (Typ noLoc) defs ctx
      check (trm_body self_var name_var) (Typ noLoc) defs body_ctx
      return $ Typ noLoc
    Ann trm_loc trm_done trm_expr trm_type -> do
      if trm_done
      then return trm_type
      else do
        check trm_expr trm_type defs ctx
        return trm_type
    _ -> throwError $ CheckErr noLoc ctx "Can't infer type"

checkExpr :: Expr -> Defs -> Except CheckErr Bool
checkExpr (Expr n typ trm) mod = do
  --traceM $ "checking: " ++ T.unpack n
  --traceM $ "type: " ++ show typ
  --traceM $ "term: " ++ show trm
  check trm typ mod []

checkDefs :: Defs -> [Except CheckErr Bool]
checkDefs mod = fmap (\(n,x) -> checkExpr x mod) (M.toList $ _defs mod)