{-# LANGUAGE TypeFamilies, KindSignatures, EmptyDataDecls, MultiParamTypeClasses, FlexibleInstances #-}
module Example2 where

import Data.Maybe
import Data.List
import Data.Map(Map)
import Data.IntMap(IntMap)
import Data.Set(Set)
import Data.IntSet(IntSet)
import Data.Sequence(Seq, (<|), (|>), (><))
import qualified Data.Map as Map
import qualified Data.IntMap as IntMap
import qualified Data.Set as Set
import qualified Data.IntSet as IntSet
import qualified Data.Sequence as Seq
import Data.IORef
import System.IO
import Control.Monad.State


---------------------------------------------------------------------
-- Library code
---------------------------------------------------------------------

data family Vis i :: * -> *
data family Inh i :: * -> *
data family Syn i :: * -> *
data family St i st :: * -> *

data Tree i s = Tree { mkSem :: IO (Sem i s) }
data Sem i s = Sem { inh :: Inh i s, syn :: Syn i s, visits :: Vis i s  }


class Decor i s where
  mkInhAttrs :: IO (Inh i s)
  mkSynAttrs :: IO (Syn i s)

mkTree :: Decor i s => (Inh i s -> Syn i s -> IO (Vis i s)) -> Tree i s
mkTree f
  = Tree $ do inh <- mkInhAttrs
              syn <- mkSynAttrs
              vis <- f inh syn
              return (Sem inh syn vis)

invoke :: Sem i s -> (Vis i s -> IO ()) -> IO ()
invoke sem f = f (visits sem)

assign :: IORef a -> a -> IO ()
assign = writeIORef

deref :: IORef a -> IO a
deref = readIORef

mkEmptyAttr :: IO (IORef a)
mkEmptyAttr = newIORef (error "uninitialized")


--
-- General data types
--

type Ident = String

type Errs = Seq Error
data Error
  = ErrMissing !Ident
  | ErrDup !Ident
  | ErrMatch !Ty !Ty
  | ErrCycle !Int !Ty
  deriving Show

data Ty
  = TyVar !Int
  | TyInt
  | TyArr !Ty !Ty
  | TyForall !IntSet !Ty
  deriving Show

newtype Subst = Subst (IntMap (Maybe Ty))

emptySubst :: Subst
emptySubst = Subst IntMap.empty

freshTy :: Subst -> (Ty, Subst)
freshTy (Subst subst)
  = let index  = IntMap.size subst
        subst' = IntMap.insert index Nothing subst
        ty     = TyVar index
    in (ty, Subst subst')

infix 4 |=>
class Substitutable a where
  (|=>) :: Subst -> a -> a

instance Substitutable Subst where
  (Subst s1) |=> (Subst s2) = Subst (s1 `IntMap.union` (IntMap.map (Subst s1 |=>) s2))

instance Substitutable a => Substitutable (Maybe a) where
  s |=> (Just x) = Just (s |=> x)
  s |=> Nothing  = Nothing

instance Substitutable Ty where
  (Subst s) |=> t@(TyVar i)
    = case IntMap.lookup i s of
        Just (Just t') -> t'
        _              -> t
  s |=> (TyArr l r)     = TyArr (s |=> l) (s |=> r)
  (Subst s) |=> (TyForall i ty) = TyForall i (Subst s |=> ty)
  _ |=> ty = ty

ftv :: Ty -> IntSet
ftv (TyVar i) = IntSet.singleton i
ftv (TyArr l r) = (ftv l) `IntSet.union` (ftv r)
ftv (TyForall i t) = ftv t `IntSet.difference` i
ftv _ = IntSet.empty

inst :: Subst -> Ty -> (Ty, Subst)
inst (Subst s) (TyForall set t)
  = let src = IntSet.toList set
        dest = map (Just . TyVar) [IntMap.size s ..]
        srep = Subst (IntMap.fromList (zip src dest))
    in ((srep |=> t), srep |=> Subst s)
inst s t = (t, s)

gen :: [Ty] -> Ty -> Ty
gen tys ty
  | IntSet.size vs > 0 = TyForall vs ty
  | otherwise          = ty
  where
    vs = ftv ty `IntSet.difference` IntSet.unions (map ftv tys)

data UniState = UniState { errs_Uni :: Errs, subst_Uni :: Subst }
type Uni a = State UniState a

unify :: Subst -> Ty -> Ty -> (Errs, Subst)
unify subst t1 t2
  = (errs, subst')
  where (UniState errs subst') = execState (unify' t1 t2) (UniState Seq.empty subst)

unify' :: Ty -> Ty -> Uni ()
unify' t1 t2
  = do s <- gets subst_Uni
       u (s |=> t1) (s |=> t2)
  where
    u (TyVar i1) (TyVar i2) | i1 == i2 = return ()
    u (TyVar i) t = bind i t
    u t (TyVar i) = bind i t
    u (TyArr t1 t2) (TyArr t3 t4)
      = do u t1 t2
           unify' t3 t4
    u _ _ = modify (\st -> st { errs_Uni = errs_Uni st |> ErrMatch t1 t2 })
    
    bind i t
      | i `IntSet.member` ftv t = modify (\st -> st { errs_Uni = errs_Uni st |> ErrCycle i t })
      | otherwise = modify (\st -> st { subst_Uni = (Subst $ IntMap.singleton i $ Just t)|=> subst_Uni st })


---------------------------------------------------------------------
-- Type analysis on main AST
---------------------------------------------------------------------

data Ast
data instance Inh Ast s = Inh_Ast {}
data instance Syn Ast s = Syn_Ast { gathName_Ast :: !(IORef (Tree (Name Ty) s)) }
data instance Vis Ast s = Vis_Ast { vName :: !(IO ())
                                  }

instance Decor Ast s where
  mkInhAttrs = do -- ...
                  return $ Inh_Ast {}
  mkSynAttrs = do tree <- mkEmptyAttr
                  return $ Syn_Ast tree

-- todo


---------------------------------------------------------------------
-- Name analysis component
---------------------------------------------------------------------

data Name a
data instance Inh (Name a) s = Inh_Name { finDefEnv_Name  :: !(IORef (Map Ident [a])) }
data instance Syn (Name a) s = Syn_Name { gathDefEnv_Name :: !(IORef (Map Ident [a])) }
data instance Vis (Name a) s = Vis_Name { vGather_Name    :: !(IO ())
                                        , vDisperse_Name  :: !(IO ()) }

instance Decor (Name a) s where
  mkInhAttrs = do fin <- mkEmptyAttr
                  return (Inh_Name fin)
  mkSynAttrs = do gath <- mkEmptyAttr
                  return (Syn_Name gath)


mk_Name_Empty :: Tree (Name a) s
mk_Name_Empty
  = mkTree $ \myInh mySyn ->
      return $ Vis_Name {
          vGather_Name   = assign (gathDefEnv_Name mySyn) Map.empty
        , vDisperse_Name = return ()
        }

mk_Name_Split :: Tree (Name a) s -> Tree (Name a) s -> Tree (Name a) s
mk_Name_Split leftTree rightTree
  = mkTree $ \myInh mySyn ->
      do left  <- mkSem leftTree
         right <- mkSem rightTree
         return $ Vis_Name {
             vGather_Name =
               do invoke left vGather_Name
                  invoke right vGather_Name

                  lGathDef <- deref $ gathDefEnv_Name $ syn left
                  rGathDef <- deref $ gathDefEnv_Name $ syn right
                  let gathDef = Map.unionWith (++) lGathDef rGathDef

                  assign (gathDefEnv_Name mySyn) gathDef
           , vDisperse_Name =
               do fin <- deref $ finDefEnv_Name myInh

                  assign (finDefEnv_Name $ inh left) fin
                  assign (finDefEnv_Name $ inh right) fin
                  invoke left vDisperse_Name
                  invoke right vDisperse_Name
           }

mk_Name_Scope :: Tree (Name a) s -> Tree (Name a) s
mk_Name_Scope subtree
  = mkTree $ \myInh mySyn ->
      do sub <- mkSem subtree
         return $ Vis_Name {
             vGather_Name =
               do invoke sub vGather_Name
                  assign (gathDefEnv_Name mySyn) Map.empty
           , vDisperse_Name =
               do fin <- deref $ finDefEnv_Name $ myInh
                  gath <- deref $ gathDefEnv_Name $ syn sub
                  assign (finDefEnv_Name $ inh sub) (Map.union gath fin)
                  invoke sub vDisperse_Name
           }

mk_Name_Use :: Tree (NameUse a) s -> Tree (Name a) s
mk_Name_Use extTree
  = mkTree $ \myInh mySyn ->
      do ext <- mkSem extTree
         return $ Vis_Name {
             vGather_Name =
               assign (gathDefEnv_Name mySyn) Map.empty
           , vDisperse_Name =
               do invoke ext vSupply_NameUse
                  name <- deref $ name_NameUse $ syn ext
                  fin  <- deref $ finDefEnv_Name myInh
                  let mbValue = maybe Nothing (Just . head) (Map.lookup name fin)
                      errs = maybe (Seq.singleton $ ErrMissing name) (const Seq.empty) mbValue
                  assign (errs_NameUse $ inh ext) errs
                  invoke ext vOutcome_NameUse             
                  case mbValue of
                    Nothing -> return ()
                    Just x  -> do assign (value_NameUse $ inh ext) x
                                  invoke ext vValue_NameUse
           }

mk_Name_Def :: Tree (NameDef a) s -> Tree (Name a) s
mk_Name_Def extTree
  = mkTree $ \myInh mySyn ->
      do ext     <- mkSem extTree
         locName <- mkEmptyAttr
         return $ Vis_Name {
             vGather_Name =
               do invoke ext vSupply_NameDef
                  name  <- deref $ name_NameDef  $ syn ext
                  value <- deref $ value_NameDef $ syn ext
                  assign (gathDefEnv_Name mySyn) (Map.singleton name [value])
                  assign locName name
           , vDisperse_Name =
               do name <- deref locName
                  fin  <- deref $ finDefEnv_Name myInh
                  let errs = case Map.lookup name fin of
                               Just xs -> if length xs == 1
                                          then Seq.empty
                                          else Seq.singleton (ErrDup name)
                  assign (errs_NameDef $ inh ext) errs
                  invoke ext vOutcome_NameDef
           }


-- External interface

data NameUse a
data instance Inh (NameUse a) s = Inh_NameUse { errs_NameUse :: IORef Errs, value_NameUse :: IORef a }
data instance Syn (NameUse a) s = Syn_NameUse { name_NameUse :: IORef Ident }
data instance Vis (NameUse a) s = Vis_NameUse { vSupply_NameUse  :: !(IO ())
                                              , vOutcome_NameUse :: !(IO ())
                                              , vValue_NameUse   :: !(IO ()) }

instance Decor (NameUse a) s where
  mkInhAttrs = do errs <- mkEmptyAttr
                  value <- mkEmptyAttr
                  return $ Inh_NameUse errs value
  mkSynAttrs = do name <- mkEmptyAttr
                  return $ Syn_NameUse name

data NameDef a
data instance Inh (NameDef a) s = Inh_NameDef { errs_NameDef :: IORef Errs }
data instance Syn (NameDef a) s = Syn_NameDef { name_NameDef :: IORef Ident, value_NameDef :: IORef a }
data instance Vis (NameDef a) s = Vis_NameDef { vSupply_NameDef  :: !(IO ())
                                              , vOutcome_NameDef :: !(IO ()) }

instance Decor (NameDef a) s where
  mkInhAttrs = do errs <- mkEmptyAttr
                  return $ Inh_NameDef errs
  mkSynAttrs = do name  <- mkEmptyAttr
                  value <- mkEmptyAttr
                  return $ Syn_NameDef name value

-- Root of the name analysis tree

data NameRoot
data instance Inh NameRoot s = Inh_NameRoot {}
data instance Syn NameRoot s = Syn_NameRoot {}
data instance Vis NameRoot s = Vis_NameRoot { visit_NameRoot :: !(IO ()) }

instance Decor NameRoot s where
  mkInhAttrs = return $ Inh_NameRoot {}
  mkSynAttrs = return $ Syn_NameRoot {}

mk_NameRoot_Root :: Tree (Name a) s -> Tree NameRoot s
mk_NameRoot_Root nameTree
  = mkTree $ \myInh mySyn ->
      do sem <- mkSem nameTree
         return $ Vis_NameRoot {
             visit_NameRoot =
               do invoke sem vGather_Name
                  gath <- deref $ gathDefEnv_Name $ syn sem
                  assign (finDefEnv_Name $ inh sem) gath
                  invoke sem vDisperse_Name
           }


---------------------------------------------------------------------
-- Type inference component
---------------------------------------------------------------------

data TypeInf
data instance Inh TypeInf s = Inh_TypeInf {}
data instance Syn TypeInf s = Syn_TypeInf { gathName_TypeInf :: !(IORef (Tree (Name Ty) s)) }
data instance Vis TypeInf s = Vis_TypeInf { vName_TypeInf    :: !(IO ())
                                          , vOutcome_TypeInf :: !(IO ())
                                          }

instance Decor TypeInf s where
  mkInhAttrs = do -- ...
                  return $ Inh_TypeInf {}
  mkSynAttrs = do tree <- mkEmptyAttr
                  return $ Syn_TypeInf tree

mk_TypeInf_Var :: Ident -> Tree TypeInf s
mk_TypeInf_Var ident
  = mkTree $ \myInh mySyn ->
      do locNmErrs  <- mkEmptyAttr
         locIdentTp <- mkEmptyAttr
         return $ Vis_TypeInf {
            vName_TypeInf =
              do let extUse = mkTree $ \myInh' mySyn' ->
                                return $ Vis_NameUse {
                                    vSupply_NameUse =
                                      assign (name_NameUse mySyn') ident
                                  , vOutcome_NameUse =
                                      do errs <- deref (errs_NameUse myInh')
                                         assign locNmErrs errs
                                  , vValue_NameUse =
                                      assign locIdentTp (value_NameUse myInh')
                                  }
                     nmTree = mk_Name_Use extUse
                 assign (gathName_TypeInf mySyn) nmTree
          , vOutcome_TypeInf =
              return ()
          }

mk_Type_App :: Tree TypeInf s -> Tree TypeInf s -> Tree TypeInf s
mk_Type_App funTree argTree
  = mkTree $ \myInh mySyn ->
      do fun <- mkSem funTree
         arg <- mkSem argTree
         return $ Vis_TypeInf {
             vName_TypeInf =
               do invoke fun vName_TypeInf
                  invoke arg vName_TypeInf
                  funNmTree <- deref (gathName_TypeInf $ syn fun)
                  argNmTree <- deref (gathName_TypeInf $ syn arg)
                  let nmTree = mk_Name_Split funNmTree argNmTree
                  assign (gathName_TypeInf mySyn) nmTree
           }

mk_TypeInf_Lam :: Ident -> Tree TypeInf s -> Tree TypeInf s
mk_TypeInf_Lam ident bodyTree
  = mkTree $ \myInh mySyn ->
      do body <- mkSem bodyTree
         locErrs <- mkEmptyAttr
         identTy <- undefined
         return $ Vis_TypeInf {
             vName_TypeInf =
               do invoke body vName_TypeInf
                  bodyNmTree <- deref (gathName_TypeInf $ syn body)
                  let extDef = mkTree $ \myInh' mySyn' ->
                                 return $ Vis_NameDef {
                                     vSupply_NameDef  =
                                       do assign (name_NameDef mySyn') ident
                                          assign (value_NameDef mySyn') identTy
                                   , vOutcome_NameDef =
                                       do errs <- deref $ errs_NameDef myInh'
                                          assign locErrs errs
                                   }
                      nmTree = mk_Name_Scope $ mk_Name_Split (mk_Name_Def extDef) bodyNmTree
                  assign (gathName_TypeInf mySyn) nmTree
           }

{-
mk_TypeInf_Let :: Ident -> Tree TypeInf s -> Tree TypeInf s -> Tree TypeInf s
mk_TypeInf_Let ident exprTree bodyTree
  = mkTree $ \myInh mySyn ->
      do expr <- mkSem exprTree
         body <- mkSem bodyTree
         return $ Vis_TypeInf {
             vName_TypeInf =
               return ()
               {-
               do invoke expr vName_TypeInf
                  invoke body vName_TypeInf
                  let extDef = mkTree $ \myInh' mySyn' ->
                                 return $ Vis_NameDef {
                                   }
               -}
           }
-}

{-
mk_Type_Fix :: Tree Type s -> Tree Type s
mk_Type_Fix expr
  = mkTree $ \myInh mySyn ->
      do sub <- mkSem subtree
         return $ Vis_Type {
             vName_Type =
               do
           }
-}

-- Root of the Type tree
data TypeRoot


---------------------------------------------------------------------
-- Kind inference component
---------------------------------------------------------------------

-- todo