package.yaml

dependencies:
- base
- mtl
- exceptions
- MonadRandom
- containers
- unification-fd
- hashable

default-extensions:
- FlexibleInstances
- MultiParamTypeClasses
- FunctionalDependencies
- UndecidableInstances
- FlexibleContexts
- DerivingStrategies
- DeriveAnyClass
- PatternSynonyms
- OverloadedStrings
- DeriveFunctor
- DeriveFoldable
- DeriveTraversable

library:
  exposed-modules:
  - Type

Type.hs

module Type where

import Control.Monad
import Control.Monad.Random
import Control.Monad.Catch
import Control.Monad.State
import Control.Monad.Except
import Control.Unification
import Control.Unification.Types

import Data.List ((\\))
import Data.String
import Data.Hashable
import Data.Foldable
import Data.Traversable
import Data.Map (Map)
import qualified Data.Map as Map
import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import Data.Set (Set)
import qualified Data.Set as Set

import Debug.Trace

data Name = Name { raw :: String, idx :: Int }
  deriving stock (Eq, Ord)

instance Show Name where
  show (Name r 0)   = r
  show (Name r idx) = r ++ "-" ++ show idx

instance IsString Name where
  fromString s = Name s 0

data Forall ty = Forall [Name] ty

instance Show a => Show (Forall a) where
  show (Forall names ty) = "/\\ " ++ unwords (map show names) ++ ". " ++ show ty

data Ty_ ty
  = C   Name ty
  | App ty ty
  | Arr ty ty
  | Star
  deriving stock (Functor, Foldable, Traversable)

type Ty = UTerm Ty_ Name

pair = UTerm $ C (Name ","    0) (star :-> star :-> star)
star = UTerm Star

instance {-# OVERLAPPING #-} Show Ty where
  show  (UVar n)         = show n
  show  (X    n)         = show n
  show ((f :$: x) :$: y) = "(" ++ show x ++ " " ++ show f ++ " " ++ show y ++ ")"
  show  (f :$: x)        = "(" ++ show f ++ " " ++ show x ++ ")"
  show  (a :-> b)        = "(" ++ show a ++ " -> " ++ show b ++ ")"
  show  (UTerm Star)     = "*"

instance Show (Ty_ Ty) where
  show = show . UTerm

pattern (:->) :: Ty -> Ty -> Ty
pattern (:*:) :: Ty -> Ty -> Ty
pattern (:$:)  :: Ty -> Ty -> Ty
pattern a :-> b =  UTerm (Arr a b)
pattern a :*: b <- UTerm (UTerm (App (UTerm (C (Name "," 0) _)) a) `App` b)
pattern a :$: b =  UTerm (App a b)

a  *: b = UTerm (UTerm (App pair a) `App` b)

pattern V s = UVar s
pattern X s <- UTerm (C s _)

data UnificationError
  = Error (UFailure Ty_ Name)
  | Unifying Ty Ty UnificationError
  | Undefined Name
  | WhileTypechecking Program UnificationError
  | WhileTypecheckingAlt (Name, [Name], Program) UnificationError
  | WhileCheckingKind Ty UnificationError
  deriving anyclass Exception

instance Show UnificationError where
  show (Error OccursFailure   {}) = "The type contains itself."
  show (Error (MismatchFailure a b)) = "The types cannot be unified: " ++ show a ++ " ~ " ++ show b
  show (Undefined name)           = "The name `" ++ show name
                                 ++ "' is not defined."
  show (Unifying a b e) =
    show e ++ "\n  While inifying: " ++ show a ++ " ~ " ++ show b

  show (WhileTypechecking p e) =
    show e ++ "\n  While typechecking: " ++ show p

  show (WhileTypecheckingAlt p e) =
    show e ++ "\n  While typechecking case alt: " ++ show p

  show (WhileCheckingKind ty e) =
    show e ++ "\n  While kind-checking: " ++ show ty

instance Fallible Ty_ Name UnificationError where
  occursFailure   name ty = Error (occursFailure   name ty)
  mismatchFailure a b     = Error (mismatchFailure a b)

instance Unifiable Ty_ where
  zipMatch (C n ty1) (C m ty2)
    | n == m = do
      return (C n (Left ty1))

  zipMatch (App f x) (App g y) = do
    return (App (Right (f, g)) (Right (x, y)))

  zipMatch (Arr a b) (Arr c d) = do
    return (Arr (Right (a, c)) (Right (b, d)))

  zipMatch Star Star = do
    return Star

  zipMatch _ _ = Nothing

instance Variable Name where
  getVarID (Name s i) = hash s + i

type M = StateT (Int, Map Name Ty) IO

runM :: M a -> IO (a, (Int, Map Name Ty))
runM ma = runStateT ma (0, Map.empty)

instance BindingMonad Ty_ Name M where
  lookupVar var = do
    (_, vars) <- get
    return (Map.lookup var vars)

  freeVar = do
    (i, vars) <- get
    put (i + 1, vars)
    return (Name "z" i)

  bindVar var term = do
    modify $ \(i, vars) -> (i, Map.insert var term vars)

data Program
  = PV Name
  | PC String Ty
  | PL Name Ty Program
  | PLet [(Name, Ty, Program)] Program
  | PTy   Ty
  | PCtor Ty
  | PApp  Program Program
  | PCase Program [(Name, [Name], Program)]
  | PHasTy Program Ty
  deriving stock Show

decorate mkE ma = do
  ma `catch` (throwM . mkE)

rethrow :: ExceptT UnificationError M a -> M a
rethrow ma = do
  res <- runExceptT ma
  either throwM return res

kindOf :: Ty -> M Ty
kindOf ty =
  decorate (WhileCheckingKind ty) $ do
    case ty of
      UVar name -> do
        maybe (throwM $ Undefined name) return
          =<< lookupVar name

      (f :$: x) -> do
        fK <- kindOf f
        xK <- kindOf x
        k  <- freeVar
        rethrow $ unify fK (xK :-> UVar k)

      UTerm (C _ k) -> do
        return k

      UTerm Star -> do
        return (UTerm Star)

typeOf :: Program -> M Ty
typeOf p = do
  traceShowM ("???? ", p)
  r <- decorate (WhileTypechecking p) $ do
    case p of
      PV name -> do
        maybe (throwM $ Undefined name) return
          =<< lookupVar name

      PC _ ty -> do
        return ty

      PL x ty body -> do
        bindVar x ty
        bodyT <- typeOf body
        return (UVar x :-> bodyT)

      PLet bindings p -> do
        for_ bindings $ \(name, ty, _) -> do
          bindVar name ty

        for_ bindings $ \(name, ty, body) -> do
          bodyT <- typeOf body
          rethrow $ unify (UVar name) bodyT

        typeOf p

      PApp f x -> do
        fT <- typeOf f
        xT <- typeOf x
        res <- freeVar
        rethrow $ unify (xT :-> UVar res) fT
        return (UVar res)

      PTy ty -> do
        kindOf ty

      PCtor ty -> do
        kind <- kindOf ty
        rethrow $ unify kind star
        return ty

      PCase o alts -> do
        oT <- typeOf o
        altTs <- for alts $ \alt @(ctor, args, body) ->
          decorate (WhileTypecheckingAlt alt) $ do
            ctorT <- lookupVar ctor
            case ctorT of
              Just ctorT -> do
                let ctorT2 = foldr (:->) oT (map UVar args)
                rethrow $ unify ctorT ctorT2
                typeOf body

              Nothing -> do
                throwM (Undefined ctor)

        res <- freeVar
        rethrow $ for altTs $ unify (UVar res)
        return (UVar res)

      PHasTy p ty -> do
        ty' <- typeOf p
        rethrow $ unify ty ty'
        return ty'

  traceShowM ("TYPE ", p, " = ", r)
  return r

test =
  PLet
    [ ("Bool", star, PTy star)
    , ("Int" , star, PTy star)
    , ("True",  UTerm (C "Bool" (UVar "Bool")), PCtor (UTerm (C "Bool" (UVar "Bool"))))
    , ("False", UTerm (C "Bool" (UVar "Bool")), PCtor (UTerm (C "Bool" (UVar "Bool"))))
    , ( "and"
      , UVar "$and"
      , PL "a" (UVar "d")
      $ PL "b" (UVar "d")
      $ (PHasTy (PV "True") (UVar "d"))
      )
    ]
    (PV "and" `PApp` PV "True")

-- data Subst = Subst (Map Name (Forall Ty))
--   deriving stock Show

-- data Bind = Name :- Forall Ty
--   deriving stock Show

-- name -: pty = do
--   when (name `elem` freeVars pty) $ do
--     throwM Cyclic

--   return (name :- pty)

-- addBind ss @(Subst old) bind = do
--   bind'@(name :- pty) <- subst ss bind
--   Subst binds <- case name `lookupTy` ss of
--     Just already -> do
--       Subst ss' <- unify pty already
--       return (Subst (Map.insert name pty ss'))

--     Nothing -> do
--       sub name pty

--   return (Subst (Map.union old binds))

-- noSubst = Subst Map.empty

-- sub n pty = do
--   _ <- n -: pty
--   return (Subst (Map.singleton n pty))

-- lookupTy name (Subst map) = Map.lookup name map

-- class HasFreeVars a where
--   freeVars :: a -> Set Name

-- class (Fresh m, MonadCatch m, Show a) => Substitutes a m where
--   subst :: Subst -> a -> m a

-- class Monad m => Fresh m where
--   fresh :: Name -> m Name

-- instance HasFreeVars Ty where
--   freeVars (Var n  ) = Set.singleton n
--   freeVars (C   _  ) = Set.empty
--   freeVars (App f x) = freeVars f <> freeVars x

-- instance HasFreeVars a => HasFreeVars (Forall a) where
--   freeVars (Poly n pty) = Set.delete n (freeVars pty)
--   freeVars (Mono    ty) = freeVars ty

-- instance (Fresh m, MonadCatch m) => Substitutes Ty m where
--   subst ss (Var name) = do
--     case name `lookupTy` ss of
--       Just poly -> instantiate poly
--       Nothing   -> return (Var name)

--   subst _ (C n) = do
--     return (C n)

--   subst ss (App f x) = do
--     f' <- subst ss f
--     x' <- subst ss x
--     return (App f' x')

-- instantiate (Poly name pty) = do
--   (_, pty') <- alphaRename (name, pty)
--   instantiate pty'

-- instantiate (Mono ty) = return ty

-- generalise
--   ::
--     ( Fresh m
--     , MonadCatch m
--     , HasFreeVars a
--     , Substitutes a m
--     )
--   => a
--   -> m (Forall a)
-- generalise ty = do
--   let fv = freeVars ty
--   foldM addParam (Mono ty) fv
--   where
--     addParam body name = do
--       -- traceShowM ("addParam", name, body)
--       (name', body') <- alphaRename (name, body)
--       return (Poly name' body')

-- alphaRename :: Substitutes a m => (Name, a) -> m (Name, a)
-- alphaRename (name, pty) = do
--   name'      <- fresh name
--   ss         <- sub name (Mono (Var name'))
--   pty'       <- subst ss pty
--   return (name', pty')

-- instance
--   ( Fresh m
--   , MonadCatch m
--   , Substitutes a m
--   ) =>
--     Substitutes (Forall a) m
--   where
--     subst ss (Poly name ty) = do
--       (name', ty') <- alphaRename (name, ty)
--       ty''         <- subst ss ty'
--       return (Poly name' ty'')

--     subst ss (Mono ty) = do
--       ty' <- subst ss ty
--       return (Mono ty')

-- otherThan name (name' :- _) = name /= name'

-- instance (Fresh m, MonadCatch m) => Substitutes Bind m where
--   subst ss (name :- pty) = do
--     pty' <- subst ss pty
--     name -: pty'

-- instance (Fresh m, MonadCatch m) => Substitutes Subst m where
--   subst ss (Subst binds) = do
--     foldM addBind ss (map (uncurry (:-)) (Map.toList binds))

-- data CannotUnify = CannotUnify
--   deriving stock    Show
--   deriving anyclass Exception

-- data WhileUnifying = WhileUnifying (Forall Ty) (Forall Ty) SomeException
--   deriving stock    Show
--   deriving anyclass Exception

-- data Cyclic = Cyclic
--   deriving stock    Show
--   deriving anyclass Exception

-- reporting ma = do
--   do { ma; return [] } `catches`
--     [ Handler $ \CannotUnify -> do
--         return ["Cannot unify."]

--     , Handler $ \(WhileUnifying a b e) -> do
--         initial <- reporting (throwM e)
--         let line = "While unifying: " ++ show a ++ " ~ " ++ show b
--         return (initial ++ [line])

--     , Handler $ \Cyclic -> do
--         return ["Type is infinite."]
--     ]

-- report ma = (mapM_ putStrLn . take 4) =<< reporting ma

-- whileUnifying a b ma = do
--   ma `catch` \e -> throwM (WhileUnifying a b e)

-- unify
--   ::
--     ( Fresh m
--     , MonadCatch m
--     )
--   => Forall Ty
--   -> Forall Ty
--   -> m Subst
-- unify left right = do
--   -- traceShowM (left, "~", right)
--   whileUnifying left right $ do
--     case (left, right) of
--       (Mono (Var n), Mono (Var m))
--         | n == m -> do
--           return noSubst

--       (Mono (Var n), pty) -> do
--         sub n pty

--       (Mono (C n), Mono (C m))
--         | n == m -> do
--           return noSubst

--       (pty, Mono (Var m)) -> do
--         sub m pty

--       (Mono (App f x), Mono (App g y)) -> do
--         ss1 <- unify (Mono f) (Mono g)
--         ss2 <- unify (Mono x) (Mono y)
--         ss  <- subst ss1 ss2
--         return ss

--       (Poly name pty1, pty2) -> do
--         pty1' <- instantiate pty1
--         unify (Mono pty1') pty2

--       (pty1, Poly name pty2) -> do
--         pty2' <- instantiate pty2
--         unify pty1 (Mono pty2')

--       _ -> throwM CannotUnify

-- instance Fresh IO where
--   fresh (Name raw _) = do
--     r <- getRandom
--     return (Name raw (abs r `rem` 100))

-- swap :: (a, b) -> (b, a)
-- swap = error ""

-- wat :: ((Int, c), (b, c)) -> (c, c)
-- wat = error ""

-- main :: IO ()
-- main = do
--   report $ do
--     idTy <- generalise ((v "a" *: v "b") --> (v "b" *: v "a"))
--     let
--       ty2 =
--         ((c "Int" *: v "c") *: (v "b" *: v "c")) --> ((v "d") *: (v "e"))

--     ss <- unify idTy (Mono ty2)
--     print ss

--     ty2' <- subst ss ty2
--     print ty2'

--     ty2'' <- generalise ty2'
--     print (freeVars ty2'')

--     print ty2''