{-# OPTIONS -w #-}

module Lambdabot.Plugin.Haskell.Free.Type where

import Control.Monad
import Lambdabot.Plugin.Haskell.Free.Parse
import Data.List
import Lambdabot.Plugin.Haskell.Free.Util
import Prelude hiding ((<>))

type TyVar = String
type TyName = String

data Type
    = TyForall TyVar Type
    | TyArr Type Type
    | TyTuple [Type]
    | TyCons TyName [Type]
    | TyVar TyVar
        deriving (Type -> Type -> Bool
(Type -> Type -> Bool) -> (Type -> Type -> Bool) -> Eq Type
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Type -> Type -> Bool
$c/= :: Type -> Type -> Bool
== :: Type -> Type -> Bool
$c== :: Type -> Type -> Bool
Eq, Int -> Type -> ShowS
[Type] -> ShowS
Type -> String
(Int -> Type -> ShowS)
-> (Type -> String) -> ([Type] -> ShowS) -> Show Type
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Type] -> ShowS
$cshowList :: [Type] -> ShowS
show :: Type -> String
$cshow :: Type -> String
showsPrec :: Int -> Type -> ShowS
$cshowsPrec :: Int -> Type -> ShowS
Show)

precTYAPP, precARROW :: Int
precTYAPP :: Int
precTYAPP = 11
precARROW :: Int
precARROW = 10

instance Pretty Type where
    prettyP :: Int -> Type -> Doc
prettyP p :: Int
p (TyForall v :: String
v t :: Type
t)
        = Bool -> Doc -> Doc
prettyParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> 0) (
            String -> Doc
text "forall" Doc -> Doc -> Doc
<+> String -> Doc
text String
v Doc -> Doc -> Doc
<> String -> Doc
text "." Doc -> Doc -> Doc
<+> Int -> Type -> Doc
forall a. Pretty a => Int -> a -> Doc
prettyP 0 Type
t
        )
    prettyP p :: Int
p (TyArr t1 :: Type
t1 t2 :: Type
t2)
        = Bool -> Doc -> Doc
prettyParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
precARROW) (
            Int -> Type -> Doc
forall a. Pretty a => Int -> a -> Doc
prettyP (Int
precARROWInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) Type
t1 Doc -> Doc -> Doc
<+> String -> Doc
text "->" Doc -> Doc -> Doc
<+> Int -> Type -> Doc
forall a. Pretty a => Int -> a -> Doc
prettyP Int
precARROW Type
t2
        )
    prettyP _ (TyTuple [])
        = Doc -> Doc
parens Doc
empty
    prettyP _ (TyTuple (t :: Type
t:ts :: [Type]
ts))
        = Doc -> Doc
parens (Int -> Type -> Doc
forall a. Pretty a => Int -> a -> Doc
prettyP 0 Type
t Doc -> Doc -> Doc
<> Int -> Doc -> [Type] -> Doc
prettyTs 0 (String -> Doc
text ",") [Type]
ts)
    prettyP _ (TyCons "[]" [t :: Type
t])
        = Doc
lbrack Doc -> Doc -> Doc
<> Int -> Type -> Doc
forall a. Pretty a => Int -> a -> Doc
prettyP 0 Type
t Doc -> Doc -> Doc
<> Doc
rbrack
    prettyP p :: Int
p (TyCons cons :: String
cons ts :: [Type]
ts)
        = Bool -> Doc -> Doc
prettyParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
precTYAPP) (
            String -> Doc
text String
cons Doc -> Doc -> Doc
<> Int -> Doc -> [Type] -> Doc
prettyTs (Int
precTYAPPInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) Doc
empty [Type]
ts
        )
    prettyP _ (TyVar v :: String
v)
        = String -> Doc
text String
v

prettyTs :: Int -> Doc -> [Type] -> Doc
prettyTs :: Int -> Doc -> [Type] -> Doc
prettyTs p :: Int
p c :: Doc
c [] = Doc
empty
prettyTs p :: Int
p c :: Doc
c (t :: Type
t:ts :: [Type]
ts) = Doc
c Doc -> Doc -> Doc
<+> Int -> Type -> Doc
forall a. Pretty a => Int -> a -> Doc
prettyP Int
p Type
t Doc -> Doc -> Doc
<> Int -> Doc -> [Type] -> Doc
prettyTs Int
p Doc
c [Type]
ts


parseType :: ParseS Type
parseType :: ParseS Type
parseType
    = ParseS Type
parseType' ParseS Type -> (Type -> ParseS Type) -> ParseS Type
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> ParseS Type) -> (Type -> Type) -> Type -> ParseS Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Type
normaliseType

parseType' :: ParseS Type
parseType' :: ParseS Type
parseType'
    = do
        Maybe Token
t <- ParseS (Maybe Token)
peekToken
        case Maybe Token
t of
            Just IdForall -> ParseS (Maybe Token)
getToken ParseS (Maybe Token) -> ParseS Type -> ParseS Type
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ParseS Type
parseForall
            _             -> ParseS Type
parseArrType
    where
        parseForall :: ParseS Type
parseForall
            = do
                Maybe Token
t <- ParseS (Maybe Token)
getToken
                case Maybe Token
t of
                    Just (QVarId v :: String
v)
                        -> ParseS Type
parseForall ParseS Type -> (Type -> ParseS Type) -> ParseS Type
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \t :: Type
t -> Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> Type -> Type
TyForall String
v Type
t)
                    Just (QVarSym ".")
                        -> ParseS Type
parseType'
                    _   -> String -> ParseS Type
forall (m :: * -> *) a. MonadFail m => String -> m a
fail "Expected variable or '.'"

        parseArrType :: ParseS Type
parseArrType
            = do
                Type
t1 <- ParseS Type
parseBType
                Maybe Token
t <- ParseS (Maybe Token)
peekToken
                case Maybe Token
t of
                    Just OpArrow
                        -> ParseS (Maybe Token)
getToken ParseS (Maybe Token) -> ParseS Type -> ParseS Type
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ParseS Type
parseType' ParseS Type -> (Type -> ParseS Type) -> ParseS Type
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \t2 :: Type
t2 ->
                            Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> Type -> Type
TyArr Type
t1 Type
t2)
                    _   -> Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return Type
t1

        parseBType :: ParseS Type
parseBType
            = do
                Type
t1 <- ParseS Type
parseAType
                case Type
t1 of
                    TyCons c :: String
c ts :: [Type]
ts
                        -> do
                            [Type]
ts' <- ParseS [Type]
parseBTypes
                            Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> [Type] -> Type
TyCons String
c ([Type]
ts[Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++[Type]
ts'))
                    _   -> Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return Type
t1

        parseBTypes :: ParseS [Type]
parseBTypes
            = (ParseS Type
parseBType ParseS Type -> (Type -> ParseS [Type]) -> ParseS [Type]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \t :: Type
t -> ParseS [Type]
parseBTypes ParseS [Type] -> ([Type] -> ParseS [Type]) -> ParseS [Type]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ts :: [Type]
ts -> [Type] -> ParseS [Type]
forall (m :: * -> *) a. Monad m => a -> m a
return (Type
tType -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:[Type]
ts))
                ParseS [Type] -> ParseS [Type] -> ParseS [Type]
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` [Type] -> ParseS [Type]
forall (m :: * -> *) a. Monad m => a -> m a
return []

        parseAType :: ParseS Type
parseAType
            = ParseS Type
parseQTyCon ParseS Type -> ParseS Type -> ParseS Type
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` ParseS Type
parseOtherAType

        parseQTyCon :: ParseS Type
parseQTyCon
            = do
                Maybe Token
t <- ParseS (Maybe Token)
getToken
                case Maybe Token
t of
                    Just OpenParen
                        -> do
                            Maybe Token
t <- ParseS (Maybe Token)
getToken
                            case Maybe Token
t of
                                Just CloseParen
                                    -> Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> [Type] -> Type
TyCons "()" [])
                                Just OpArrow
                                    -> Token -> ParseS ()
match Token
CloseParen
                                        ParseS () -> ParseS Type -> ParseS Type
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> [Type] -> Type
TyCons "->" [])
                                Just Comma
                                    -> Int -> ParseS Type
parseQTyConTuple 1
                                _   -> String -> ParseS Type
forall (m :: * -> *) a. MonadFail m => String -> m a
fail "Badly formed type constructor"
                    Just OpenBracket
                        -> Token -> ParseS ()
match Token
CloseBracket ParseS () -> ParseS Type -> ParseS Type
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> [Type] -> Type
TyCons "[]" [])
                    Just (QConId v :: String
v)
                        -> Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> [Type] -> Type
TyCons String
v [])
                    _   -> String -> ParseS Type
forall (m :: * -> *) a. MonadFail m => String -> m a
fail "Badly formed type constructor"

        parseQTyConTuple :: Int -> ParseS Type
        parseQTyConTuple :: Int -> ParseS Type
parseQTyConTuple i :: Int
i
            = do
                Maybe Token
t <- ParseS (Maybe Token)
getToken
                case Maybe Token
t of
                    Just Comma
                        -> Int -> ParseS Type
parseQTyConTuple (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1)
                    Just CloseParen
                        -> Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> [Type] -> Type
TyCons ("(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> ShowS
forall a. Int -> [a] -> [a]
take Int
i (Char -> String
forall a. a -> [a]
repeat ',') String -> ShowS
forall a. [a] -> [a] -> [a]
++ ")") [])
                    _   -> String -> ParseS Type
forall (m :: * -> *) a. MonadFail m => String -> m a
fail "Badly formed type constructor"

        parseOtherAType :: ParseS Type
parseOtherAType
            = do
                Maybe Token
t1 <- ParseS (Maybe Token)
getToken
                case Maybe Token
t1 of
                    Just OpenParen
                        -> do
                            Type
t <- ParseS Type
parseType'
                            [Type] -> ParseS Type
parseTuple [Type
t]
                    Just OpenBracket
                        -> ParseS Type
parseType' ParseS Type -> (Type -> ParseS Type) -> ParseS Type
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \t :: Type
t -> Token -> ParseS ()
match Token
CloseBracket
                                        ParseS () -> ParseS Type -> ParseS Type
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> [Type] -> Type
TyCons "[]" [Type
t])
                    Just (QVarId v :: String
v)
                        -> Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> Type
TyVar String
v)
                    _   -> String -> ParseS Type
forall (m :: * -> *) a. MonadFail m => String -> m a
fail "Badly formed type"

        parseTuple :: [Type] -> ParseS Type
parseTuple ts :: [Type]
ts
            = do
                Maybe Token
t1 <- ParseS (Maybe Token)
getToken
                case Maybe Token
t1 of
                    Just CloseParen
                        -> case [Type]
ts of
                                [t :: Type
t] -> Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return Type
t
                                _   -> Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return ([Type] -> Type
TyTuple ([Type] -> [Type]
forall a. [a] -> [a]
reverse [Type]
ts))
                    Just Comma
                        -> do
                            Type
t <- ParseS Type
parseType'
                            [Type] -> ParseS Type
parseTuple (Type
tType -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:[Type]
ts)

normaliseType :: Type -> Type
normaliseType :: Type -> Type
normaliseType t :: Type
t
    = let (fvs :: [String]
fvs,nt :: Type
nt) = Type -> ([String], Type)
normaliseType' Type
t
      in (String -> Type -> Type) -> Type -> [String] -> Type
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr String -> Type -> Type
TyForall Type
nt ([String] -> [String]
forall a. Eq a => [a] -> [a]
nub [String]
fvs)
    where
        normaliseType' :: Type -> ([String], Type)
normaliseType' t :: Type
t@(TyVar v :: String
v)
            = ([String
v],Type
t)
        normaliseType' (TyForall v :: String
v t' :: Type
t')
            = let (fvs :: [String]
fvs,t :: Type
t) = Type -> ([String], Type)
normaliseType' Type
t'
              in ((String -> Bool) -> [String] -> [String]
forall a. (a -> Bool) -> [a] -> [a]
filter (String -> String -> Bool
forall a. Eq a => a -> a -> Bool
/=String
v) [String]
fvs, String -> Type -> Type
TyForall String
v Type
t)
        normaliseType' (TyArr t1 :: Type
t1 t2 :: Type
t2)
            = let
                (fvs1 :: [String]
fvs1,t1' :: Type
t1') = Type -> ([String], Type)
normaliseType' Type
t1
                (fvs2 :: [String]
fvs2,t2' :: Type
t2') = Type -> ([String], Type)
normaliseType' Type
t2
              in
                ([String]
fvs1[String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++[String]
fvs2, Type -> Type -> Type
TyArr Type
t1' Type
t2')
        normaliseType' (TyTuple ts :: [Type]
ts)
            = let
                fvsts :: [([String], Type)]
fvsts = (Type -> ([String], Type)) -> [Type] -> [([String], Type)]
forall a b. (a -> b) -> [a] -> [b]
map Type -> ([String], Type)
normaliseType' [Type]
ts
                fvs :: [String]
fvs = [[String]] -> [String]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ((([String], Type) -> [String]) -> [([String], Type)] -> [[String]]
forall a b. (a -> b) -> [a] -> [b]
map ([String], Type) -> [String]
forall a b. (a, b) -> a
fst [([String], Type)]
fvsts)
                ts' :: [Type]
ts' = (([String], Type) -> Type) -> [([String], Type)] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map ([String], Type) -> Type
forall a b. (a, b) -> b
snd [([String], Type)]
fvsts
              in ([String]
fvs, [Type] -> Type
TyTuple [Type]
ts')
        normaliseType' (TyCons c :: String
c ts :: [Type]
ts)
            = let
                fvsts :: [([String], Type)]
fvsts = (Type -> ([String], Type)) -> [Type] -> [([String], Type)]
forall a b. (a -> b) -> [a] -> [b]
map Type -> ([String], Type)
normaliseType' [Type]
ts
                fvs :: [String]
fvs = [[String]] -> [String]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ((([String], Type) -> [String]) -> [([String], Type)] -> [[String]]
forall a b. (a -> b) -> [a] -> [b]
map ([String], Type) -> [String]
forall a b. (a, b) -> a
fst [([String], Type)]
fvsts)
                ts' :: [Type]
ts' = (([String], Type) -> Type) -> [([String], Type)] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map ([String], Type) -> Type
forall a b. (a, b) -> b
snd [([String], Type)]
fvsts
              in case String
c of
                    "->" -> case [Type]
ts' of
                        [t1 :: Type
t1,t2 :: Type
t2] -> ([String]
fvs, Type -> Type -> Type
TyArr Type
t1 Type
t2)
                        _ -> String -> ([String], Type)
forall a. HasCallStack => String -> a
error "Arrow type should have 2 arguments"
                    _ -> case String -> Maybe Int
forall a. Num a => String -> Maybe a
checkTuple String
c of
                        Just i :: Int
i
                            -> if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
ts'
                                then ([String]
fvs, [Type] -> Type
TyTuple [Type]
ts')
                                else String -> ([String], Type)
forall a. HasCallStack => String -> a
error "Tuple type has the wrong number of arguments"
                        Nothing
                            -> ([String]
fvs, String -> [Type] -> Type
TyCons String
c [Type]
ts')

        checkTuple :: String -> Maybe a
checkTuple ('(':')':cs :: String
cs)
            = a -> Maybe a
forall a. a -> Maybe a
Just 0
        checkTuple ('(':cs :: String
cs)
            = a -> String -> Maybe a
forall t. Num t => t -> String -> Maybe t
checkTuple' 1 String
cs
        checkTuple _
            = Maybe a
forall a. Maybe a
Nothing

        checkTuple' :: t -> String -> Maybe t
checkTuple' k :: t
k ")"
            = t -> Maybe t
forall a. a -> Maybe a
Just t
k
        checkTuple' k :: t
k (',':cs :: String
cs)
            = t -> String -> Maybe t
checkTuple' (t
kt -> t -> t
forall a. Num a => a -> a -> a
+1) String
cs
        checkTuple' _ _
            = Maybe t
forall a. Maybe a
Nothing

readType :: String -> Type
readType :: String -> Type
readType s :: String
s
    = case ParseS Type -> [Token] -> ParseResult Type
forall a. ParseS a -> [Token] -> ParseResult a
parse ParseS Type
parseType (String -> [Token]
lexer String
s) of
        ParseSuccess t :: Type
t [] -> Type
t
        ParseSuccess t :: Type
t _  -> String -> Type
forall a. HasCallStack => String -> a
error "Extra stuff at end of type"
        ParseError msg :: String
msg    -> String -> Type
forall a. HasCallStack => String -> a
error String
msg

-- vim: ts=4:sts=4:expandtab:ai