content: -- Copyright 2023 Google LLC
  --
  -- Use of this source code is governed by a BSD-style
  -- license that can be found in the LICENSE file or at
  -- https://developers.google.com/open-source/licenses/bsd
  
  module QueryTypePure where
  
  import Types.Primitives
  import Types.Core
  import Types.Top
  import IRVariants
  import Name
  
  class HasType (r::IR) (e::E) | e -> r where
    getType :: e n -> Type r n
  
  class HasEffects (e::E) (r::IR) | e -> r where
    getEffects :: e n -> EffectRow r n
  
  getTyCon :: HasType SimpIR e => e n -> TyCon SimpIR n
  getTyCon e = con where TyCon con = getType e
  
  isPure :: (IRRep r, HasEffects e r) => e n -> Bool
  isPure e = case getEffects e of
    Pure -> True
    _    -> False
  
  -- === querying types implementation ===
  
  instance IRRep r => HasType r (AtomBinding r) where
    getType = \case
      LetBound    (DeclBinding _ e)  -> getType e
      MiscBound   ty                 -> ty
      SolverBound (InfVarBound ty)   -> ty
      SolverBound (SkolemBound ty)   -> ty
      SolverBound (DictBound   ty)   -> ty
      NoinlineFun ty _               -> ty
      TopDataBound e                 -> getType e
      FFIFunBound piTy _ -> TyCon $ Pi piTy
  
  litType :: LitVal -> BaseType
  litType v = case v of
    Int64Lit   _ -> Scalar Int64Type
    Int32Lit   _ -> Scalar Int32Type
    Word8Lit   _ -> Scalar Word8Type
    Word32Lit  _ -> Scalar Word32Type
    Word64Lit  _ -> Scalar Word64Type
    Float64Lit _ -> Scalar Float64Type
    Float32Lit _ -> Scalar Float32Type
    PtrLit ty _  -> PtrType ty
  
  typeBinOp :: BinOp -> BaseType -> BaseType
  typeBinOp binop xTy = case binop of
    IAdd   -> xTy;  ISub   -> xTy
    IMul   -> xTy;  IDiv   -> xTy
    IRem   -> xTy;
    ICmp _ -> Scalar Word8Type
    FAdd   -> xTy;  FSub   -> xTy
    FMul   -> xTy;  FDiv   -> xTy;
    FPow   -> xTy
    FCmp _ -> Scalar Word8Type
    BAnd   -> xTy;  BOr    -> xTy
    BXor   -> xTy
    BShL   -> xTy;  BShR   -> xTy
  
  typeUnOp :: UnOp -> BaseType -> BaseType
  typeUnOp = const id  -- All unary ops preserve the type of the input
  
  
  instance IRRep r => HasType r (AtomVar r) where
    getType (AtomVar _ ty) = ty
    {-# INLINE getType #-}
  
  instance IRRep r => HasType r (Atom r) where
    getType = \case
      Stuck t _ -> t
      Con e -> getType e
  
  instance HasType CoreIR (Dict CoreIR) where
    getType = \case
      StuckDict t _ -> t
      DictCon e -> getType e
  
  instance HasType CoreIR (DictCon CoreIR) where
    getType = \case
      InstanceDict t _ _ -> t
      DataData t -> toType $ DataDictType t
      IxFin n -> toType $ IxDictType (FinTy n)
      IxRawFin _ -> toType $ IxDictType IdxRepTy
  
  instance HasType CoreIR CType where
    getType = \case
      TyCon _ -> TyKind
      StuckTy t _ -> t
  
  instance HasType CoreIR NewtypeTyCon where
    getType _ = TyKind
  
  getNewtypeType :: NewtypeCon n -> CType n
  getNewtypeType con = case con of
    NatCon              -> TyCon $ NewtypeTyCon Nat
    FinCon n            -> TyCon $ NewtypeTyCon $ Fin n
    UserADTData sn d xs -> TyCon $ NewtypeTyCon $ UserADTType sn d xs
  
  instance IRRep r => HasType r (Con r) where
    getType = \case
      Lit l          -> toType $ BaseType $ litType l
      ProdCon xs     -> toType $ ProdType $ map getType xs
      SumCon tys _ _ -> toType $ SumType tys
      HeapVal        -> toType HeapType
      Lam (CoreLamExpr piTy _) -> toType $ Pi piTy
      DepPair _ _ ty -> toType $ DepPairTy ty
      Eff _ -> EffKind
      DictConAtom d -> getType d
      NewtypeCon con _ -> getNewtypeType con
      TyConAtom _ -> TyKind
  
  getSuperclassType :: RNest CBinder n l -> Nest CBinder l l' -> Int -> CType n
  getSuperclassType _ Empty = error "bad index"
  getSuperclassType bsAbove (Nest b@(_:>t) bs) = \case
    0 -> ignoreHoistFailure $ hoist bsAbove t
    i -> getSuperclassType (RNest bsAbove b) bs (i-1)
  
  instance IRRep r => HasType r (Expr r) where
    getType expr = case expr of
      App (EffTy _ ty) _ _ -> ty
      TopApp (EffTy _ ty) _ _ -> ty
      TabApp t _ _ -> t
      Atom x   -> getType x
      Block (EffTy _ ty) _ -> ty
      TabCon _ ty _ -> ty
      PrimOp op -> getType op
      Case _ _ (EffTy _ resultTy) -> resultTy
      ApplyMethod (EffTy _ t) _ _ _ -> t
      Project t _ _ -> t
      Unwrap t _ -> t
  
  instance HasType SimpIR RepVal where
    getType (RepVal ty _) = ty
  
  instance IRRep r => HasType r (DAMOp r) where
    getType = \case
      AllocDest ty -> RawRefTy ty
      Place _ _ -> UnitTy
      Freeze ref -> case getType ref of
        RawRefTy ty -> ty
        ty -> error $ "Not a reference type: " ++ show ty
      Seq _ _ _ cinit _ -> getType cinit
      RememberDest _ d _ -> getType d
  
  instance IRRep r => HasType r (PrimOp r) where
    getType primOp = case primOp of
      BinOp op x _ -> TyCon $ BaseType $ typeBinOp op $ getTypeBaseType x
      UnOp  op x   -> TyCon $ BaseType $ typeUnOp  op $ getTypeBaseType x
      Hof  (TypedHof (EffTy _ ty) _) -> ty
      MemOp op -> getType op
      MiscOp op -> getType op
      VectorOp op -> getType op
      DAMOp           op -> getType op
      RefOp ref m -> case getType ref of
        TyCon (RefType _ s) -> case m of
          MGet        -> s
          MPut _      -> UnitTy
          MAsk        -> s
          MExtend _ _ -> UnitTy
          IndexRef t _ -> t
          ProjRef t _ -> t
        _ -> error "not a reference type"
  
  getTypeBaseType :: (IRRep r, HasType r e) => e n -> BaseType
  getTypeBaseType e = case getType e of
    TyCon (BaseType b) -> b
    ty -> error $ "Expected a base type. Got: " ++ show ty
  
  instance IRRep r => HasType r (MemOp r) where
    getType = \case
      IOAlloc _ -> PtrTy (CPU, Scalar Word8Type)
      IOFree _ -> UnitTy
      PtrOffset arr _ -> getType arr
      PtrLoad ptr -> do
        let PtrTy (_, t) = getType ptr
        toType $ BaseType t
      PtrStore _ _ -> UnitTy
  
  instance IRRep r => HasType r (VectorOp r) where
    getType = \case
      VectorBroadcast _ vty -> vty
      VectorIota vty -> vty
      VectorIdx _ _ vty -> vty
      VectorSubref ref _ vty -> case getType ref of
        TyCon (RefType h _) -> TyCon $ RefType h vty
        ty -> error $ "Not a reference type: " ++ show ty
  
  instance IRRep r => HasType r (MiscOp r) where
    getType = \case
      Select _ x _ -> getType x
      ThrowError t     -> t
      ThrowException t -> t
      CastOp t _       -> t
      BitcastOp t _    -> t
      UnsafeCoerce t _ -> t
      GarbageVal t -> t
      SumTag _     -> TagRepTy
      ToEnum t _   -> t
      OutputStream -> toType $ BaseType $ hostPtrTy $ Scalar Word8Type
        where hostPtrTy ty = PtrType (CPU, ty)
      ShowAny _ -> rawStrType -- TODO: constrain `ShowAny` to have `HasCore r`
      ShowScalar _ -> toType $ ProdType [IdxRepTy, rawFinTabType (IdxRepVal showStringBufferSize) CharRepTy]
  
  rawStrType :: IRRep r => Type r n
  rawStrType = case newName "n" of
    Abs b v -> do
      let tabTy = rawFinTabType (toAtom $ AtomVar v IdxRepTy) CharRepTy
      TyCon $ DepPairTy $ DepPairType ExplicitDepPair (b:>IdxRepTy) tabTy
  
  -- `n` argument is IdxRepVal, not Nat
  rawFinTabType :: IRRep r => Atom r n -> Type r n -> Type r n
  rawFinTabType n eltTy = IxType IdxRepTy (DictCon (IxRawFin n)) ==> eltTy
  
  tabIxType :: TabPiType r n -> IxType r n
  tabIxType (TabPiType d (_:>t) _) = IxType t d
  
  typesAsBinderNest
    :: (SinkableE e, HoistableE e, IRRep r)
    => [Type r n] -> e n -> Abs (Nest (Binder r)) e n
  typesAsBinderNest types body = toConstBinderNest types body
  
  nonDepPiType :: [CType n] -> EffectRow CoreIR n -> CType n -> CorePiType n
  nonDepPiType argTys eff resultTy = case typesAsBinderNest argTys (PairE eff resultTy) of
    Abs bs (PairE eff' resultTy') -> do
      let expls = nestToList (const Explicit) bs
      CorePiType ExplicitApp expls bs $ EffTy eff' resultTy'
  
  nonDepTabPiType :: IRRep r => IxType r n -> Type r n -> TabPiType r n
  nonDepTabPiType (IxType t d) resultTy =
    case toConstAbsPure resultTy of
      Abs b resultTy' -> TabPiType d (b:>t) resultTy'
  
  corePiTypeToPiType :: CorePiType n -> PiType CoreIR n
  corePiTypeToPiType (CorePiType _ _ bs effTy) = PiType bs effTy
  
  coreLamToTopLam :: CoreLamExpr n -> TopLam CoreIR n
  coreLamToTopLam (CoreLamExpr ty f) = TopLam False (corePiTypeToPiType ty) f
  
  (==>) :: IRRep r => IxType r n -> Type r n -> Type r n
  a ==> b = TyCon $ TabPi $ nonDepTabPiType a b
  
  litFinIxTy :: Int -> IxType SimpIR n
  litFinIxTy n = finIxTy $ IdxRepVal $ fromIntegral n
  
  finIxTy :: Atom SimpIR n -> IxType SimpIR n
  finIxTy n = IxType IdxRepTy (DictCon (IxRawFin n))
  
  -- === querying effects implementation ===
  
  instance IRRep r => HasEffects (Expr r) r where
    getEffects = \case
      Atom _ -> Pure
      Block (EffTy eff _) _ -> eff
      App (EffTy eff _) _ _ -> eff
      TopApp (EffTy eff _) _ _ -> eff
      TabApp _ _ _ -> Pure
      Case _ _ (EffTy effs _) -> effs
      TabCon _ _ _      -> Pure
      ApplyMethod (EffTy eff _) _ _ _ -> eff
      PrimOp primOp -> getEffects primOp
      Project _ _ _ -> Pure
      Unwrap _ _ -> Pure
  
  instance IRRep r => HasEffects (DeclBinding r) r where
    getEffects (DeclBinding _ expr) = getEffects expr
    {-# INLINE getEffects #-}
  
  instance IRRep r => HasEffects (PrimOp r) r where
    getEffects = \case
      UnOp  _ _   -> Pure
      BinOp _ _ _ -> Pure
      VectorOp _  -> Pure
      MemOp op -> case op of
        IOAlloc  _    -> OneEffect IOEffect
        IOFree   _    -> OneEffect IOEffect
        PtrLoad  _    -> OneEffect IOEffect
        PtrStore _ _  -> OneEffect IOEffect
        PtrOffset _ _ -> Pure
      MiscOp op -> case op of
        ThrowException _ -> OneEffect ExceptionEffect
        Select _ _ _     -> Pure
        ThrowError _     -> Pure
        CastOp _ _       -> Pure
        UnsafeCoerce _ _ -> Pure
        GarbageVal _     -> Pure
        BitcastOp _ _    -> Pure
        SumTag _         -> Pure
        ToEnum _ _       -> Pure
        OutputStream     -> Pure
        ShowAny _        -> Pure
        ShowScalar _     -> Pure
      RefOp ref m -> case getType ref of
        TyCon (RefType h _) -> case m of
          MGet      -> OneEffect (RWSEffect State  h)
          MPut    _ -> OneEffect (RWSEffect State  h)
          MAsk      -> OneEffect (RWSEffect Reader h)
          -- XXX: We don't verify the base monoid. See note about RunWriter.
          MExtend _ _ -> OneEffect (RWSEffect Writer h)
          IndexRef _ _ -> Pure
          ProjRef _ _  -> Pure
        _ -> error "not a ref"
      DAMOp op -> case op of
        Place    _ _  -> OneEffect InitEffect
        Seq eff _ _ _ _        -> eff
        RememberDest eff _ _ -> eff
        AllocDest _ -> Pure -- is this correct?
        Freeze _    -> Pure -- is this correct?
      Hof (TypedHof (EffTy eff _) _) -> eff
    {-# INLINE getEffects #-}
  