{-# 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 = Int
11
precARROW :: Int
precARROW = Int
10

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

prettyTs :: Int -> Doc -> [Type] -> Doc
prettyTs :: Int -> Doc -> [Type] -> Doc
prettyTs Int
p Doc
c [] = Doc
empty
prettyTs Int
p Doc
c (Type
t:[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 Token
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
            Maybe Token
_             -> ParseS Type
parseArrType
    where
        parseForall :: ParseS Type
parseForall
            = do
                Maybe Token
t <- ParseS (Maybe Token)
getToken
                case Maybe Token
t of
                    Just (QVarId 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
>>= \Type
t -> Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> Type -> Type
TyForall String
v Type
t)
                    Just (QVarSym String
".")
                        -> ParseS Type
parseType'
                    Maybe Token
_   -> String -> ParseS Type
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"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 Token
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
>>= \Type
t2 ->
                            Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> Type -> Type
TyArr Type
t1 Type
t2)
                    Maybe Token
_   -> 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 String
c [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
_   -> 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
>>= \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
>>= \[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 Token
OpenParen
                        -> do
                            Maybe Token
t <- ParseS (Maybe Token)
getToken
                            case Maybe Token
t of
                                Just Token
CloseParen
                                    -> Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> [Type] -> Type
TyCons String
"()" [])
                                Just Token
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 String
"->" [])
                                Just Token
Comma
                                    -> Int -> ParseS Type
parseQTyConTuple Int
1
                                Maybe Token
_   -> String -> ParseS Type
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Badly formed type constructor"
                    Just Token
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 String
"[]" [])
                    Just (QConId String
v)
                        -> Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> [Type] -> Type
TyCons String
v [])
                    Maybe Token
_   -> String -> ParseS Type
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Badly formed type constructor"

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

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

        parseTuple :: [Type] -> ParseS Type
parseTuple [Type]
ts
            = do
                Maybe Token
t1 <- ParseS (Maybe Token)
getToken
                case Maybe Token
t1 of
                    Just Token
CloseParen
                        -> case [Type]
ts of
                                [Type
t] -> Type -> ParseS Type
forall (m :: * -> *) a. Monad m => a -> m a
return Type
t
                                [Type]
_   -> 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 Token
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 Type
t
    = let ([String]
fvs,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 String
v)
            = ([String
v],Type
t)
        normaliseType' (TyForall String
v Type
t')
            = let ([String]
fvs,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 Type
t1 Type
t2)
            = let
                ([String]
fvs1,Type
t1') = Type -> ([String], Type)
normaliseType' Type
t1
                ([String]
fvs2,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 [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 String
c [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
                    String
"->" -> case [Type]
ts' of
                        [Type
t1,Type
t2] -> ([String]
fvs, Type -> Type -> Type
TyArr Type
t1 Type
t2)
                        [Type]
_ -> String -> ([String], Type)
forall a. HasCallStack => String -> a
error String
"Arrow type should have 2 arguments"
                    String
_ -> case String -> Maybe Int
forall a. Num a => String -> Maybe a
checkTuple String
c of
                        Just 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 String
"Tuple type has the wrong number of arguments"
                        Maybe Int
Nothing
                            -> ([String]
fvs, String -> [Type] -> Type
TyCons String
c [Type]
ts')

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

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

readType :: String -> Type
readType :: String -> Type
readType 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 Type
t [] -> Type
t
        ParseSuccess Type
t [Token]
_  -> String -> Type
forall a. HasCallStack => String -> a
error String
"Extra stuff at end of type"
        ParseError String
msg    -> String -> Type
forall a. HasCallStack => String -> a
error String
msg

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