{-# OPTIONS -fglasgow-exts #-}

-- A module for modelling all of the derivations discussed in the draft paper, 
-- "Simulating Dependent Types with Guarded Algebraic Datatypes" by Jim Apple
-- and Wes Weimer.

module Lift where

import Monad(liftM)

{- We restrict contexts to hold kinds and types for strings, so as to 
   disallow a conflict between a lifted type's kind as inferred and the
   kind listed in the context -}
data Context = Empty
             | Kinded String Kind Context
             | Typed  String RestrictedConstructorType Context deriving (Show)

data ConstructorType = ConstructorType [(Variable,Kind)] [Type] Data [Type] deriving (Show)
data RestrictedConstructorType = RestrictedConstructorType [(Variable,Kind)] [Type] String [Type] deriving (Show)


data Constructor = Constructor String
                 | LiftConstructor Constructor deriving (Show)

data Data = Data String
          | LiftData Data
          | Tag Constructor deriving (Show)

data Kind = Star
          | Kind `KindArrow` Kind deriving (Show)

data Type = Type `TypeArrow` Type
          | Type `TypeApply` Type
          | Forall Variable Kind Type
          | Lambda Variable Kind Type
          | Name Data
          | Var Variable deriving (Show)

data Variable = Variable String deriving (Show)

infixl `TypeApply`
infixr `KindArrow`
infixr `TypeArrow`

arity Star = 0
arity (_ `KindArrow` k) = 1 + arity k

liftKind :: Kind -> Kind
liftKind Star = Star `KindArrow` Star
liftKind (k1 `KindArrow` k2) = (liftKind k1) `KindArrow` (liftKind k2)

typeOf :: Context -> Constructor -> Maybe ConstructorType
-- the first three are the usual context-searching for names
typeOf Empty (Constructor s) = Nothing
typeOf (Typed s1 (RestrictedConstructorType vk t0 s2 t1) gamma) (Constructor s3) = 
    if s1 == s3
    then Just  (ConstructorType vk t0 (Data s2) t1)
    else typeOf gamma (Constructor s3)
typeOf (Kinded _ _ gamma) (Constructor s) = typeOf gamma (Constructor s)
typeOf gamma (LiftConstructor c) = 
    do 
      --This must be a constructor available in the context
      ConstructorType vks ts d rs <- typeOf gamma c
      --We lift all of the arguments
      t's <- sequence $ map liftName ts
      r's <- sequence $ map liftName rs
      return $
             let 
                 -- The new variables we quantify over. We make sure to pull them 
                 -- from the available free names
                 usVar = take (length ts) (freeIn $ (names gamma) ++ (names vks) ++ (names t's) ++ (names d) ++ (names r's))
                 us = map Var usVar
                 -- We're going to keep the v's but lift the k's
                 (vs,ks) = unzip vks
             in 
               ConstructorType 
               ((zip usVar (repeat Star)) ++ (zip vs (map liftKind ks)))
               (zipWith TypeApply t's us)
               (LiftData d)
               (r's ++ [foldl TypeApply (Name $ Tag c) us])

-- Which of v0, v1, v2 . . . v314 . . .  are not mentioned in the argument
freeIn :: [String] -> [Variable]
freeIn ss = map (Variable) (filter (not . (flip elem ss)) (map ('v':) (map show [0 ..])))

{- Helper function for getting used variables from different parts
   of the grammar -}
class Names a where
    names :: a -> [String]
instance Names Context where
    names Empty = []
    names (Kinded s _ gamma) = s : (names gamma)
    names (Typed  s _ gamma) = s : (names gamma)
instance Names Variable where
    names (Variable s) = [s]
instance (Names a, Names b) => Names (a,b) where
    names (x,y) = (names x)++(names y)
instance Names Kind where
    names _ = []
instance Names Type where
    names (t1 `TypeArrow` t2) = (names t1)++(names t2)
    names (t1 `TypeApply` t2) = (names t1)++(names t2)
    names (Forall v k t) = (names v)++(names k)++(names t)
    names (Lambda v k t) = (names v)++(names k)++(names t)
    names (Name d) = names d
    names (Var v) = names v
instance Names Data where
    names (Data s) = [s]
    names (LiftData d) = names d
    names (Tag c) = names c
instance Names Constructor where
    names (Constructor s) = [s]
    names (LiftConstructor c) = names c
instance Names a => Names [a] where
    names x = concatMap names x

liftName :: Type -> Maybe Type
liftName (t1 `TypeApply` r1) = 
    do t2 <- liftName t1
       r2 <- liftName r1
       return $ t2 `TypeApply` r2
liftName (Forall v1 k t1) =
    let v2 = head $ freeIn ((names v1) ++ (names t1))
    in do 
        t2 <- liftName t1 
        return $ Lambda v2 Star $ Forall v1 (liftKind k) (t2 `TypeApply` (Var v2))
liftName (Lambda v k t) =
    liftM (Lambda v (liftKind k)) (liftName t)
liftName (Name d) = Just $ Name $ LiftData d
liftName (Var v) = Just $ Var v
liftName _ = Nothing

kindOf :: Context -> Data -> Maybe Kind
-- the first three are the usual context-searching for names
kindOf Empty (Data s) = Nothing
kindOf (Kinded s1 k gamma) (Data s2) =
    if s2 == s1
    then Just k
    else kindOf gamma (Data s2)
kindOf (Typed _ _ gamma) (Data s) = 
    kindOf gamma (Data s)
kindOf gamma (LiftData d) =
    case kindOf gamma d of
      (Just k) -> Just (liftKind k)
      _ -> Nothing
kindOf gamma (Tag c) = 
    case typeOf gamma c of
      Nothing -> Nothing
      Just (ConstructorType _ ts _ _) -> Just $ foldr KindArrow Star $ map (const Star) ts  --(makeKind (length ts))

-- Figure out whether a type admits equality.
-- The second argument is the list of type constructors we may assume do admit equality.
-- It should start out empty and grow only as recursing.
singleton :: Context -> [String] -> Type -> Bool
singleton gamma x (t1 `TypeApply` t2) = singleton gamma x t1 && singleton gamma x t2
singleton gamma x (Forall _ _ t) = singleton gamma x t
singleton gamma x (Name (Tag _)) = True
singleton gamma x (Name (LiftData d)) = singleton gamma x (Name d)
singleton gamma x (Name (Data s)) = 
    if s `elem` x then True else
        case kindOf gamma (Data s) of
          Nothing -> False
          Just _ -> let cs = restrictedConstructorsOf s gamma
                        toCheck = (concatMap (arguments . snd) cs)++(concatMap (resultArguments . snd) cs)
                    in all (singleton gamma (s:x)) toCheck
singleton gamma x (Var _) = True
singleton gamma x _ = False

restrictedConstructorsOf :: String -> Context -> [(String,RestrictedConstructorType)]
restrictedConstructorsOf s Empty = []
restrictedConstructorsOf s (Kinded _ _ gamma) = restrictedConstructorsOf s gamma
restrictedConstructorsOf s (Typed n r@(RestrictedConstructorType _ _ s' _) gamma) = 
    let rest = restrictedConstructorsOf s gamma in
    if s == s'
    then (n,r) : rest
    else rest

constructorsOf :: Data -> Context -> Maybe [(Constructor,ConstructorType)]
constructorsOf x y = constructorsOf' x y y id
constructorsOf' _ _ Empty _ = Just []
constructorsOf' (Tag _) _ _ _ = Just []
constructorsOf' d y (Kinded _ _ gamma) f = constructorsOf' d y gamma f
constructorsOf' (LiftData d) y gamma f = constructorsOf' d y gamma (f . LiftConstructor)
constructorsOf' (Data s) y (Typed n r@(RestrictedConstructorType _ _ s' _) gamma) f = 
    do 
      rest <- constructorsOf' (Data s) y gamma f
      (if s == s'
       then 
           do r <- typeOf y (f (Constructor n))
              return $ (f (Constructor n), r) : rest
       else return rest)

arguments :: RestrictedConstructorType -> [Type]
arguments (RestrictedConstructorType _ a _ _) = a
resultArguments :: RestrictedConstructorType -> [Type]
resultArguments (RestrictedConstructorType _ _ _ a) = a