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 RuntimePrint (showAny) where
  
  import Control.Monad.Reader
  import Data.Functor
  
  import Builder
  import Core
  import Err
  import IRVariants
  import MTL1
  import Name
  import CheapReduction
  import PPrint
  import Types.Core
  import Types.Source
  import Types.Primitives
  import QueryType
  import Util (enumerate)
  
  newtype Printer (n::S) (a :: *) = Printer { runPrinter' :: ReaderT1 (Atom CoreIR) (BuilderM CoreIR) n a }
          deriving ( Functor, Applicative, Monad, EnvReader, MonadReader (Atom CoreIR n)
                   , Fallible, ScopeReader, MonadFail, EnvExtender, CBuilder, ScopableBuilder CoreIR)
  type Print n = Printer n ()
  
  showAny :: EnvReader m => Atom CoreIR n -> m n (CExpr n)
  showAny x = liftPrinter $ showAnyRec (sink x)
  
  liftPrinter
    :: EnvReader m
    => (forall l. (DExt n l, Emits l) => Print l)
    -> m n (CExpr n)
  liftPrinter cont = liftBuilder $ buildBlock $ withBuffer \buf ->
    runReaderT1 buf (runPrinter' cont)
  
  getBuffer :: Printer n (CAtom n)
  getBuffer = ask
  
  emitCharTab :: Emits n => CAtom n -> Print n
  emitCharTab tab = do
    buf <- getBuffer
    extendBuffer buf tab
  
  emitChar :: Emits n => CAtom n -> Print n
  emitChar tab = do
    buf <- getBuffer
    pushBuffer buf tab
  
  emitLit :: Emits n => String -> Print n
  emitLit s = stringLitAsCharTab s >>= emitCharTab
  
  emitCharLit :: Emits n => Char -> Print n
  emitCharLit c = emitChar $ charRepVal c
  
  showAnyRec :: forall n. Emits n => CAtom n -> Print n
  showAnyRec atom = case getType atom of
    TyCon con -> showAnyTyCon con atom
    StuckTy _ e -> error $ "unexpected stuck type expression: " ++ pprint e
  
  showAnyTyCon :: forall n. Emits n => TyCon CoreIR n -> CAtom n -> Print n
  showAnyTyCon tyCon atom = case tyCon of
    BaseType bt -> case bt of
      Vector _ _ -> error "not implemented"
      PtrType _  -> printTypeOnly "pointer"
      Scalar _ -> do
        (n, tab) <- fromPair =<< emit (ShowScalar atom)
        logicalTabTy <- finTabTyCore (Con $ NewtypeCon NatCon n) CharRepTy
        tab' <- emit $ UnsafeCoerce logicalTabTy tab
        emitCharTab tab'
    -- TODO: we could do better than this but it's not urgent because raw sum types
    -- aren't user-facing.
    SumType _ -> printAsConstant
    RefType _ _ -> printTypeOnly "reference"
    HeapType    -> printAsConstant
    ProdType _ -> do
      xs <- getUnpacked atom
      parens $ sepBy ", " $ map rec xs
    -- TODO: traverse the type and print out data components
    TypeKind -> printAsConstant
    Pi _ -> printTypeOnly "function"
    TabPi _ -> brackets $ forEachTabElt atom \iOrd x -> do
      isFirst <- emit $ BinOp (ICmp Equal) iOrd (NatVal 0)
      void $ emitIf isFirst UnitTy (return UnitVal) (emitLit ", " >> return UnitVal)
      rec x
    NewtypeTyCon tc -> case tc of
      Fin _ -> rec =<< unwrapNewtype atom
      Nat -> do
        n <- unwrapNewtype atom
        -- Cast to Int so that it prints in decimal instead of hex
        let intTy = toType $ BaseType (Scalar Int64Type)
        emit (CastOp intTy n) >>= rec
      EffectRowKind    -> printAsConstant
      -- hack to print strings nicely. TODO: make `Char` a newtype
      UserADTType "List" _ (TyConParams [Explicit] [Con (TyConAtom (BaseType (Scalar (Word8Type))))]) -> do
        charTab <- applyProjections [ProjectProduct 1, UnwrapNewtype] atom
        emitCharLit '"'
        emitCharTab charTab
        emitCharLit '"'
      UserADTType tySourceName defName params -> do
        def <- lookupTyCon defName
        conDefs <- instantiateTyConDef def params
        case conDefs of
          ADTCons [con] -> showDataCon con =<< unwrapNewtype atom
          ADTCons cons -> void $ buildCase atom UnitTy \i arg -> do
            showDataCon (sink $ cons !! i) arg
            return UnitVal
          StructFields fields -> do
            emitLit $ pprint tySourceName
            parens do
              sepBy ", " $ (enumerate fields) <&> \(i, _) ->
                rec =<< projectStruct i atom
        where
          showDataCon :: Emits n' => DataConDef n' -> CAtom n' -> Print n'
          showDataCon (DataConDef sn _ _ projss) arg = do
            case projss of
              [] -> emitLit $ pprint sn
              _ -> parens do
                emitLit (pprint sn ++ " ")
                sepBy " " $ projss <&> \projs ->
                  -- we use `init` to strip off the `UnwrapCompoundNewtype` since
                  -- we're already under the case alternative
                  rec =<< applyProjections (init projs) arg
    DepPairTy _ -> parens do
      (x, y) <- fromPair atom
      rec x >> emitLit " ,> " >> rec y
    -- Done well, this could let you inspect the results of dictionary synthesis
    -- and maybe even debug synthesis failures.
    DictTy _ -> printAsConstant
    where
      rec :: Emits n' => CAtom n' -> Print n'
      rec = showAnyRec
  
      printTypeOnly :: String -> Print n
      printTypeOnly thingName = do
        ty <- return $ getType atom
        emitLit $ "<" <> thingName <> " of type " <> pprint ty <> ">"
  
      printAsConstant :: Print n
      printAsConstant = emitLit $ pprint atom
  
  parens :: Emits n => Print n -> Print n
  parens x = emitCharLit '(' >> x >> emitCharLit ')'
  
  brackets :: Emits n => Print n -> Print n
  brackets x = emitCharLit '[' >> x >> emitCharLit ']'
  
  sepBy :: forall n. Emits n => String -> [Print n] -> Print n
  sepBy s xsTop = rec xsTop where
    rec :: [Print n] -> Print n
    rec = \case
      []   -> return ()
      [x]  -> x
      x:xs -> x >> emitLit s >> rec xs
  
  -- === Builder helpers (consider moving these to Builder.hs) ===
  
  withBuffer
    :: Emits n
    => (forall l . (Emits l, DExt n l) => CAtom l -> BuilderM CoreIR l ())
    -> BuilderM CoreIR n (CAtom n)
  withBuffer cont = do
    lam <- withFreshBinder "h" (TyCon HeapType) \h -> do
      bufTy <- bufferTy (toAtom $ binderVar h)
      withFreshBinder "buf" bufTy \b -> do
        let eff = OneEffect (RWSEffect State (toAtom $ sink $ binderVar h))
        body <- buildBlock do
          cont $ sink $ toAtom $ binderVar b
          return UnitVal
        let binders = BinaryNest h b
        let expls = [Inferred Nothing Unify, Explicit]
        let piTy = CorePiType ExplicitApp expls binders $ EffTy eff UnitTy
        let lam = LamExpr (BinaryNest h b) body
        return $ toAtom $ CoreLamExpr piTy lam
    applyPreludeFunction "with_stack_internal" [lam]
  
  bufferTy :: EnvReader m => CAtom n -> m n (CType n)
  bufferTy h = do
    t <- strType
    return $ RefTy h (PairTy NatTy t)
  
  -- argument has type `Fin n => Word8`
  extendBuffer :: (Emits n, CBuilder m) => CAtom n -> CAtom n -> m n ()
  extendBuffer buf tab = do
    RefTy h _ <- return $ getType buf
    TyCon (TabPi t) <- return $ getType tab
    n <- applyIxMethodCore Size (tabIxType t) []
    void $ applyPreludeFunction "stack_extend_internal" [n, h, buf, tab]
  
  -- argument has type `Word8`
  pushBuffer :: (Emits n, CBuilder m) => CAtom n -> CAtom n -> m n ()
  pushBuffer buf x = do
    RefTy h _ <- return $ getType buf
    void $ applyPreludeFunction "stack_push_internal" [h, buf, x]
  
  stringLitAsCharTab :: (Emits n, CBuilder m) => String -> m n (CAtom n)
  stringLitAsCharTab s = do
    t <- finTabTyCore (NatVal $ fromIntegral $ length s) CharRepTy
    emit $ TabCon Nothing t (map charRepVal s)
  
  finTabTyCore :: (Fallible1 m, EnvReader m) => CAtom n -> CType n -> m n (CType n)
  finTabTyCore n eltTy = return $ IxType (FinTy n) (DictCon $ IxFin n) ==> eltTy
  
  getPreludeFunction :: EnvReader m => SourceName -> m n (CAtom n)
  getPreludeFunction sourceName = do
    lookupSourceMap sourceName >>= \case
      Just uvar -> case uvar of
        UAtomVar v -> toAtom <$> toAtomVar v
        _ -> notfound
      Nothing -> notfound
   where notfound = error $ "Function not defined: " ++ pprint sourceName
  
  applyPreludeFunction :: (Emits n, CBuilder m) => SourceName -> [CAtom n] -> m n (CAtom n)
  applyPreludeFunction name args = do
    f <- getPreludeFunction name
    naryApp f args
  
  strType :: forall n m. EnvReader m => m n (CType n)
  strType = constructPreludeType "List" $ TyConParams [Explicit] [toAtom (CharRepTy :: CType n)]
  
  constructPreludeType :: EnvReader m => SourceName -> TyConParams n -> m n (CType n)
  constructPreludeType sourceName params = do
    lookupSourceMap sourceName >>= \case
      Just uvar -> case uvar of
        UTyConVar v -> return $ toType $ UserADTType sourceName v params
        _ -> notfound
      Nothing -> notfound
   where notfound = error $ "Type constructor not defined: " ++ pprint sourceName
  
  forEachTabElt
    :: (Emits n, ScopableBuilder CoreIR m)
    => CAtom n
    -> (forall l. (Emits l, DExt n l) => CAtom l -> CAtom l -> m l ())
    -> m n ()
  forEachTabElt tab cont = do
    TyCon (TabPi t) <- return $ getType tab
    let ixTy = tabIxType t
    void $ buildFor "i" Fwd ixTy \i -> do
      x <- tabApp (sink tab) (toAtom i)
      i' <- applyIxMethodCore Ordinal (sink ixTy) [toAtom i]
      cont i' x
      return $ UnitVal
  