TiDecorate

module TiDecorate where
import BaseSyntax
import HasBaseStruct
import Substitute
import SubstituteBaseStruct(substE') --hmm
import FreeNames
import DefinedNames
import Syntax(HsPatI,HsTypeI,HsDeclI,HsExpI,Rec(..))
import qualified Syntax as S
import PrettyPrint
import PrettySymbols
import TI hiding (Subst)
import TiBase()
import TiBaseStruct(tcE,tcP,tcD)
import List(nub)
--import Maybe(mapMaybe)
--import QualNames
import MUtils(apSnd)
--import IOExts(trace)

--------------------------------------------------------------------------------
--type TiModule i = HsModuleI i (TiDecls i)
data TiDecls i = Decs [TiDecl i] (DeclsType i) deriving Show

type DStruct i = DI i (TiExp i) (TiPat i) (TiDecls i) (HsTypeI i) [HsTypeI i] (HsTypeI i)
type EStruct i = EI i (TiExp i) (TiPat i) (TiDecls i) (HsTypeI i) [HsTypeI i]
type PStruct i = PI i (TiPat i)

newtype TiDecl i = Dec (DStruct i) deriving Show

data TiExp i
  = Exp (EStruct i)
  | TiSpec (HsIdentI i) (Scheme i) [HsTypeI i] -- specialize a polymorphic identifier
  | TiTyped (TiExp i) (HsTypeI i)
  deriving Show

data TiPat i
  = Pat (PStruct i)
  | TiPApp (TiPat i) (TiPat i) -- for fromInteger & fromRational & n+k
  | TiPSpec (HsIdentI i) (Scheme i) [HsTypeI i] -- specialize a polymorphic identifier
  | TiPTyped (TiPat i) (HsTypeI i)
  deriving Show

--------------------------------------------------------------------------------
instance MapExp (TiExp i) (TiDecls i) where
  mapExp f (Decs ds dt) = Decs (mapExp f ds) dt

instance MapExp (TiExp i) (TiDecl i) where mapExp = mapExpRec

instance Subst i (TiExp i) where
  subst s e =
    case e of
      Exp be -> substE' Exp s be
      TiTyped e t -> TiTyped (subst s e) t
      TiSpec (HsVar x) sc ts ->
          case isId sx of
            Just y -> TiSpec y sc ts
	    _ -> sx -- type annotation is lost!
	where sx = s x
      _ -> e 
--------------------------------------------------------------------------------

instance (PrintableOp i,ValueId i,TypeId i,Fresh i,HasSrcLoc i,TypedId PId i)
      => TypeCheckDecl i (HsDeclI i) (TiDecls i) where
  tcDecl bs = tcD bs . struct

instance AddDeclsType i (TiDecls i) where
  addDeclsType dt1 (Decs ds dt2) = Decs ds (dt1+++dt2)
  rmDeclsType (Decs ds dt) = Decs ds ([],[])

instance (PrintableOp i,ValueId i,TypeId i,Fresh i,HasSrcLoc i,TypedId PId i)
      =>  TypeCheck i (HsExpI i) (Typed i (TiExp i)) where tc (S.Exp e) = tcE e

instance (Printable i,ValueId i,TypeId i,Fresh i)
      =>  TypeCheck i (HsPatI i) (Typed i (TiPat i)) where tc (S.Pat e) = tcP e
{-
instance GetSigs i [Pred i] (Type i) (TiDecls i) where
  getSigs (Decs ds _) = mapMaybe getSig ds
    where
      getSig (Dec (HsTypeSig s is c tp)) = Just (s,is,c,tp)
      getSig _                    = Nothing

instance (TypeVar i,ValueId i) => DeclInfo i (TiDecl i) where
  explicitlyTyped {-m-} bs = explicitlyTyped {-m-} bs . struct
  isUnrestricted b = isUnrestricted b . struct

instance DeclInfo i (TiDecls i)
-}
--------------------------------------------------------------------------------

type DeclsUseType i = ([TAssump i],[UAssump i])
type UAssump i = Typing (HsIdentI i) (Scheme i,Maybe (Type i))

instance EnvFrom ds env => EnvFrom (HsModuleI m i ds) env where
  accEnv = accEnv . hsModDecls

instance TypeVar i => EnvFrom (TiDecls i) (DeclsUseType i) where
  accEnv (Decs ds dt) = (no_use dt+++) . accEnv ds

no_use = apSnd ((map.fmap) (flip (,) Nothing))
drop_use = apSnd ((map.fmap) fst)

instance TypeVar i => EnvFrom (TiDecl i) (DeclsUseType i) where
  accEnv = accEnv . struct

instance (EnvFrom e env,EnvFrom p env,EnvFrom ds env)
      => EnvFrom (DI i e p ds t c tp) env where
  accEnv = accDI ignore accEnv accEnv accEnv ignore ignore ignore

instance TypeVar i => EnvFrom (TiPat i) (DeclsUseType i) where
  accEnv (Pat p) = accEnv p
  accEnv (TiPSpec v@(HsVar _) t as) = (([],[v:>:(t,specType t as)])+++)
  accEnv (TiPApp (TiPSpec {}) p2) = accEnv p2
  accEnv (TiPApp p1 p2) = accEnv p1 . accEnv p2
  accEnv (TiPTyped p t) = accEnv p
  accEnv _ = id

instance EnvFrom p env => EnvFrom (PI i p) env where
  accEnv = accPI ignore accEnv

instance TypeVar i => EnvFrom (TiExp i) (DeclsUseType i) where
  accEnv (Exp e) = accEnv e
  accEnv (TiSpec v@(HsVar _) t as) = (([],[v:>:(t,specType t as)])+++)
  accEnv (TiTyped e t) = accEnv e
  accEnv _ = id

specType (Forall _ kgs (_:=>t)) as =
    if map tyvar gs==as
    then Nothing
    else Just (apply s t)
  where s = S (zip gs as)
        gs = tdom kgs

instance (EnvFrom e env,EnvFrom p env,EnvFrom ds env)
      => EnvFrom (EI i e p ds t c) env where
  accEnv = accEI ignore accEnv accEnv accEnv ignore ignore

ignore = const id

--------------------------------------------------------------------------------
instance HasBaseStruct (TiDecls i) [TiDecl i] where base ds = Decs ds ([],[])
instance HasBaseStruct (TiDecl i) (DStruct i) where base    = Dec
instance HasBaseStruct (TiExp i) (EStruct i)  where base    = Exp
instance HasBaseStruct (TiPat i) (PStruct i)  where base    = Pat

instance Rec (TiDecl i) (DStruct i) where rec = base; struct (Dec d) = d
instance GetBaseStruct (TiDecl i) (DStruct i) where basestruct = Just . struct
instance GetBaseStruct (TiExp i) (EStruct i) where
  basestruct e = case e of
    Exp e' -> Just e'
    _ -> Nothing

instance HasId i (TiExp i) where
  ident = Exp . ident
  isId (Exp e) = isId e
  isId (TiSpec x _ []) = Just x
  isId _ = Nothing --hmm

instance HasId i (TiPat i) where
  ident = Pat . ident
  isId (Pat e) = isId e
  isId (TiPSpec x _ []) = Just x
  isId _ = Nothing --hmm

--instance HasLit (TiExp i) where lit = hsLit loc0
--instance HasLit (TiPat i) where lit = hsPLit loc0

instance HasCoreSyntax i (TiExp i) where
  app = hsApp
  tuple = hsTuple
  list = hsList
--paren = hsParen

instance HasCoreSyntax i (TiPat i) where
  app = TiPApp
  tuple = hsPTuple loc0 -- !!
  list = hsPList loc0 -- !!
--paren = hsParen

instance HasTypeApp i (TiExp i) where spec=TiSpec
instance HasTypeApp i (TiPat i) where spec=TiPSpec

instance HasTypeAnnot i (TiExp i) where typeAnnot=TiTyped
instance HasTypeAnnot i (TiPat i) where typeAnnot=TiPTyped

instance HasDef (TiDecls i) (TiDecl i) where
  nullDef (Decs ds _) = null ds
  noDef = Decs [] ([],[])
  consDef d (Decs ds ts) = Decs (d:ds) ts
  toDefs ds = Decs ds ([],[])
  appendDef (Decs ds1 ts1) (Decs ds2 ts2) = Decs (ds1++ds2) (ts1+++ts2)
  filterDefs p (Decs ds ts) = Decs (filter p ds) ts -- keep all type info?!

instance HasDefs (TiDecls i) (TiDecl i) where
  fromDefs (Decs ds _) = ds

instance HasSrcLoc (TiDecl i) where srcLoc (Dec d) = srcLoc d

instance HasAbstr i (TiDecl i) where abstract xs (Dec d)      = Dec (abstract xs d)
instance HasAbstr i (TiDecls i) where abstract xs (Decs ds ti) = Decs (abstract xs ds) ti
--------------------------------------------------------------------------------
instance DefinedNames i (TiDecls i) where definedNames (Decs ds ti) = definedNames ds
instance ClassMethods i (TiDecls i) where classMethods i cnt (Decs ds ti) = classMethods i cnt ds
instance MapDefinedNames i (TiDecls i) where
   mapDefinedNames f (Decs ds ti) = mapTiDeclsNames Decs f ds ti

mapTiDeclsNames decs f ds (ks,ts) =
      decs (mapDefinedNames f ds) (ks,mapVars f ts)
   where mapVars f = map $ emap $ mapHsIdent2 f id

instance DefinedNames i (TiDecl i) where definedNames = definedNamesRec
instance ClassMethods i (TiDecl i) where classMethods = classMethodsRec
instance AddName i (TiDecl i)

instance MapDefinedNames i (TiDecl i) where
   mapDefinedNames = mapDefinedNamesRec

instance DefinedNames i (TiPat i) where
  definedNames (Pat p) = definedNames p
  definedNames (TiPSpec (HsVar x) sc ts) = [value x]
  definedNames (TiPSpec _ sc ts) = []
  definedNames (TiPTyped p t) = definedNames p
  definedNames (TiPApp p1 p2) = appDef p1 [p2]
    where
      appDef (TiPApp p1 p2) ps = appDef p1 (p2:ps)
      appDef (TiPSpec c@(HsCon _) sc ts) ps = definedNames ps
      appDef (TiPSpec x sc ts) ps = [] -- overloaded literal, binds nothing

instance MapDefinedNames i (TiPat i) where
  mapDefinedNames f p =
    case p of
      Pat p -> Pat (m p)
      TiPSpec (HsVar x) sc ts -> TiPSpec (HsVar (f x)) sc ts
      TiPTyped p t -> TiPTyped (m p) t
      TiPApp p1 p2 -> mApp p1 [p2]
    where
      m x = mapDefinedNames f x

      mApp (TiPApp p1 p2) ps = mApp p1 (p2:ps)
      mApp p@(TiPSpec c@(HsCon _) _ _) ps = apps (p:m ps)
      mApp p ps = apps (p:ps) -- overloaded literal, binds nothing

      apps = foldl1 TiPApp

instance Eq i => FreeNames i (TiDecl i) where freeNames = freeNamesRec
instance Eq i => FreeNames i (TiDecls i) where freeNames (Decs ds ti) = freeNames ds
instance Eq i => FreeNames i (TiExp i) where
  freeNames (Exp e) = freeNames e
  freeNames (TiSpec x sc ts) = [val x] -- hmm
  freeNames (TiTyped e t) = freeNames e -- hmm

instance (Eq i) => FreeNames i (TiPat i) where
  freeNames (Pat p) = freeNames p
  freeNames (TiPSpec (HsCon c) sc ts) = [FreeNames.con c]
  freeNames (TiPSpec _ sc ts) = []
  freeNames (TiPTyped p t) = freeNames p -- hmm
  freeNames p@(TiPApp p1 p2) = appFree p1 [p2]
    where
      {- An application is either
           1. A construtor applied to a number of subpatterns, or
           2. fromInteger/fromRational/negate, applied to a dictionary
              and a literal. -}
      appFree (TiPApp p1 p2) ps = appFree p1 (p2:ps)
      appFree (TiPSpec (HsCon c) sc ts) ps = nub (FreeNames.con c:freeNames ps)
      appFree (TiPSpec x sc ts) ps = nub (val x:concatMap argFree ps)
	   -- x must be negate, fromInteger or fromRational
      appFree (Pat (HsPId x@(HsVar _))) [] = [] -- dictionary bound in where clause
      appFree (Pat (HsPId _)) ps = error "Bug: TiDecorate.appFree (Pat (HsPId _))"
      appFree (Pat p) ps = error ("Bug: TiDecorate.appFree (Pat _)")
      appFree p ps = error ("Bug: TiDecorate.appFree p ps")

      argFree (Pat (HsPLit _ _)) = []
      argFree p = appFree p []

--------------------------------------------------------------------------------
instance Eq i => HasLocalDef i (TiExp i) (TiDecl i) where
   --letvar x e (Dec d) = Dec (letvar x e d)
    letvar x e = mapRec (letvar x e)

instance Eq i => HasLocalDef i (TiExp i) (TiDecls i) where
   letvar x e (Decs ds ts) = Decs (letvar x e ds) ts

-- For these instances, tmap f should apply f to all types inserted during
-- type checking, since they might contain type variables that have been
-- instantiated later.
instance TypeVar i => Types i (TiDecl i) where
  tmap f (Dec d) = Dec (tmap f d)
--tv = ...

instance TypeVar i => Types i (TiDecls i) where
  tmap f (Decs ds (ks,ts)) = Decs (tmap f ds) (ks,tmap f ts)
--tv = ...

instance TypeVar i => Types i (TiExp i) where
  tmap f (Exp e) = Exp (mapEI id (tmap f) (tmap f) (tmap f) id id e)
  tmap f (TiSpec x sc ts) = TiSpec x (tmap f sc) (tmap f ts)
  tmap f (TiTyped e t) = TiTyped (tmap f e) (f t)
--tv = ...

instance TypeVar i => Types i (TiPat i) where
  tmap f (Pat p) = Pat (mapPI id (tmap f) p)
  tmap f (TiPApp p1 p2) = TiPApp (tmap f p1) (tmap f p2)
  tmap f (TiPSpec x sc ts) = TiPSpec x (tmap f sc) (tmap f ts)
  tmap f (TiPTyped e t) = TiPTyped (tmap f e) (f t)
--tv = ...
--------------------------------------------------------------------------------
instance (TypeId i,ValueId i,PrintableOp i) => Printable (TiDecl i) where
  ppi (Dec d)      = ppi d
instance (TypeId i,ValueId i,PrintableOp i) => Printable (TiDecls i) where
  ppi (Decs ds (ks,ts)) = vcat ds $$
			  ppIfTypeInfo (sep [ppi "{-",ppi ks,ppi ts,ppi "-}"])
  ppis (Decs ds ([],[])) = ppis ds
  ppis (Decs ds (ks,ts)) = ppis ds ++
                           map ppIfTypeInfo [ppi "{-",ppi ks,ppi ts,ppi "-}"]

instance (TypeId i,ValueId i,PrintableOp i) => Printable (TiExp i) where
   ppi (Exp e)      = ppi e
   ppi (TiSpec x sc []) = wrap x <> ppIfDebug ("{-"<>sc<>"-}")
   ppi (TiSpec x sc ts) = wrap x <>
                          ppIfDebug ("{-"<>sc<>"@"<>fsep (map wrap ts)<>"-}")
   ppi (TiTyped e t) = e <> ppIfDebug ("{-"<>el<>t<>"-}")

   wrap (Exp e) = wrap e
   wrap e = ppi e


instance (TypeId i,ValueId i,PrintableOp i) => Printable (TiPat i) where
   ppi (Pat e)      = ppi e
   ppi (TiPSpec x sc []) = wrap x <> ppIfDebug ("{-"<>sc<>"-}")
   ppi (TiPSpec x sc ts) = wrap x <> ppIfTypeInfo ("{-"<>ppIfDebug (sc<>"@")<>fsep (map wrap ts)<>"-}")
   ppi (TiPApp p1 p2) = fsep [ ppi p1, letNest (wrap p2) ]
   ppi (TiPTyped p t) = p<>ppIfDebug ("{-"<>el<>t<>"-}")

   wrap (Pat e) = wrap e
   wrap p@(TiPApp {}) = parens (ppi p)
   wrap p = ppi p