{-|
Copyright        : (c) Galois, Inc. 2025
Maintainer       : Langston Barrett <langston@galois.com>
-}

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TypeApplications #-}

module Lang.Crucible.Pretty
  ( IntrinsicPrettyFn(..)
  , IntrinsicPrinters(..)
  , ppRegVal
  ) where

import qualified Data.List as List
import           Data.Kind (Type)
import qualified Data.Text as Text

import qualified Prettyprinter as PP

import qualified Data.Parameterized.Context as Ctx
import           Data.Parameterized.Map (MapF)
import qualified Data.Parameterized.Map as MapF
import           Data.Parameterized.SymbolRepr (Symbol)
import           Data.Parameterized.TraversableFC (toListFC)

-- what4
import qualified What4.Interface as W4

-- crucible
import qualified Lang.Crucible.FunctionHandle as C
import qualified Lang.Crucible.Simulator.Intrinsics as C
import qualified Lang.Crucible.Simulator.RegMap as C
import qualified Lang.Crucible.Simulator.SymSequence as C
import qualified Lang.Crucible.Types as C

-- | Function for pretty-printing an intrinsic type
--
-- TODO: Upstream to Crucible
type IntrinsicPrettyFn :: Type -> Symbol -> Type
newtype IntrinsicPrettyFn sym nm
  = IntrinsicPrettyFn
    (forall ctx ann.
      Ctx.Assignment C.TypeRepr ctx ->
      C.Intrinsic sym nm ctx ->
      PP.Doc ann)

newtype IntrinsicPrinters sym
  = IntrinsicPrinters
    { forall sym.
IntrinsicPrinters sym -> MapF SymbolRepr (IntrinsicPrettyFn sym)
getIntrinsicPrinters :: MapF C.SymbolRepr (IntrinsicPrettyFn sym) }

-- TODO: Complete all the cases, upstream to Crucible
-- | Pretty-print a 'C.RegValue'.
ppRegVal ::
  forall sym tp ann.
  W4.IsExpr (W4.SymExpr sym) =>
  IntrinsicPrinters sym ->
  C.TypeRepr tp ->
  C.RegValue' sym tp ->
  PP.Doc ann
ppRegVal :: forall sym (tp :: CrucibleType) ann.
IsExpr (SymExpr sym) =>
IntrinsicPrinters sym -> TypeRepr tp -> RegValue' sym tp -> Doc ann
ppRegVal IntrinsicPrinters sym
iFns TypeRepr tp
tp (C.RV RegValue sym tp
v) =
  case TypeRepr tp
tp of
    -- Base types
    TypeRepr tp
C.BoolRepr -> SymExpr sym BaseBoolType -> Doc ann
forall (tp :: BaseType) ann. SymExpr sym tp -> Doc ann
forall (e :: BaseType -> *) (tp :: BaseType) ann.
IsExpr e =>
e tp -> Doc ann
W4.printSymExpr RegValue sym tp
SymExpr sym BaseBoolType
v
    C.BVRepr NatRepr n
_width -> SymExpr sym (BaseBVType n) -> Doc ann
forall (tp :: BaseType) ann. SymExpr sym tp -> Doc ann
forall (e :: BaseType -> *) (tp :: BaseType) ann.
IsExpr e =>
e tp -> Doc ann
W4.printSymExpr RegValue sym tp
SymExpr sym (BaseBVType n)
v
    TypeRepr tp
C.ComplexRealRepr -> SymExpr sym BaseComplexType -> Doc ann
forall (tp :: BaseType) ann. SymExpr sym tp -> Doc ann
forall (e :: BaseType -> *) (tp :: BaseType) ann.
IsExpr e =>
e tp -> Doc ann
W4.printSymExpr RegValue sym tp
SymExpr sym BaseComplexType
v
    C.FloatRepr FloatInfoRepr flt
_fi -> SymExpr sym (SymInterpretedFloatType sym flt) -> Doc ann
forall (tp :: BaseType) ann. SymExpr sym tp -> Doc ann
forall (e :: BaseType -> *) (tp :: BaseType) ann.
IsExpr e =>
e tp -> Doc ann
W4.printSymExpr RegValue sym tp
SymExpr sym (SymInterpretedFloatType sym flt)
v
    C.IEEEFloatRepr FloatPrecisionRepr ps
_fpp -> SymExpr sym (BaseFloatType ps) -> Doc ann
forall (tp :: BaseType) ann. SymExpr sym tp -> Doc ann
forall (e :: BaseType -> *) (tp :: BaseType) ann.
IsExpr e =>
e tp -> Doc ann
W4.printSymExpr RegValue sym tp
SymExpr sym (BaseFloatType ps)
v
    TypeRepr tp
C.IntegerRepr -> SymExpr sym BaseIntegerType -> Doc ann
forall (tp :: BaseType) ann. SymExpr sym tp -> Doc ann
forall (e :: BaseType -> *) (tp :: BaseType) ann.
IsExpr e =>
e tp -> Doc ann
W4.printSymExpr RegValue sym tp
SymExpr sym BaseIntegerType
v
    TypeRepr tp
C.NatRepr -> SymExpr sym BaseIntegerType -> Doc ann
forall (tp :: BaseType) ann. SymExpr sym tp -> Doc ann
forall (e :: BaseType -> *) (tp :: BaseType) ann.
IsExpr e =>
e tp -> Doc ann
W4.printSymExpr (SymNat sym -> SymExpr sym BaseIntegerType
forall sym. SymNat sym -> SymInteger sym
W4.natToIntegerPure RegValue sym tp
SymNat sym
v)
    TypeRepr tp
C.RealValRepr -> SymExpr sym BaseRealType -> Doc ann
forall (tp :: BaseType) ann. SymExpr sym tp -> Doc ann
forall (e :: BaseType -> *) (tp :: BaseType) ann.
IsExpr e =>
e tp -> Doc ann
W4.printSymExpr RegValue sym tp
SymExpr sym BaseRealType
v
    C.StringRepr StringInfoRepr si
_si -> SymExpr sym (BaseStringType si) -> Doc ann
forall (tp :: BaseType) ann. SymExpr sym tp -> Doc ann
forall (e :: BaseType -> *) (tp :: BaseType) ann.
IsExpr e =>
e tp -> Doc ann
W4.printSymExpr RegValue sym tp
SymExpr sym (BaseStringType si)
v
    C.SymbolicArrayRepr Assignment BaseTypeRepr (idx ::> tp1)
_idxs BaseTypeRepr t
_tp' -> SymExpr sym (BaseArrayType (idx ::> tp1) t) -> Doc ann
forall (tp :: BaseType) ann. SymExpr sym tp -> Doc ann
forall (e :: BaseType -> *) (tp :: BaseType) ann.
IsExpr e =>
e tp -> Doc ann
W4.printSymExpr RegValue sym tp
SymExpr sym (BaseArrayType (idx ::> tp1) t)
v
    C.SymbolicStructRepr Assignment BaseTypeRepr ctx
_tys -> SymExpr sym (BaseStructType ctx) -> Doc ann
forall (tp :: BaseType) ann. SymExpr sym tp -> Doc ann
forall (e :: BaseType -> *) (tp :: BaseType) ann.
IsExpr e =>
e tp -> Doc ann
W4.printSymExpr RegValue sym tp
SymExpr sym (BaseStructType ctx)
v

    -- Trivial cases
    TypeRepr tp
C.UnitRepr -> Doc ann
"()"
    TypeRepr tp
C.CharRepr -> Word16 -> Doc ann
forall ann. Word16 -> Doc ann
forall a ann. Pretty a => a -> Doc ann
PP.pretty Word16
RegValue sym tp
v

    -- Simple recursive cases
    C.RecursiveRepr SymbolRepr nm
symb CtxRepr ctx
tyCtx ->
      IntrinsicPrinters sym
-> TypeRepr (UnrollType nm ctx)
-> RegValue' sym (UnrollType nm ctx)
-> Doc ann
forall sym (tp :: CrucibleType) ann.
IsExpr (SymExpr sym) =>
IntrinsicPrinters sym -> TypeRepr tp -> RegValue' sym tp -> Doc ann
ppRegVal IntrinsicPrinters sym
iFns (SymbolRepr nm -> CtxRepr ctx -> TypeRepr (UnrollType nm ctx)
forall (nm :: Symbol) (ctx :: Ctx CrucibleType).
IsRecursiveType nm =>
SymbolRepr nm -> CtxRepr ctx -> TypeRepr (UnrollType nm ctx)
forall (ctx :: Ctx CrucibleType).
SymbolRepr nm -> CtxRepr ctx -> TypeRepr (UnrollType nm ctx)
C.unrollType SymbolRepr nm
symb CtxRepr ctx
tyCtx) (forall sym (tp :: CrucibleType).
RegValue sym tp -> RegValue' sym tp
C.RV @sym (RolledType sym nm ctx -> RegValue sym (UnrollType nm ctx)
forall sym (nm :: Symbol) (ctx :: Ctx CrucibleType).
RolledType sym nm ctx -> RegValue sym (UnrollType nm ctx)
C.unroll RolledType sym nm ctx
RegValue sym tp
v))
    C.SequenceRepr TypeRepr tp1
tp' ->
      (RegValue sym tp1 -> Doc ann)
-> SymSequence sym (RegValue sym tp1) -> Doc ann
forall sym a ann.
IsExpr (SymExpr sym) =>
(a -> Doc ann) -> SymSequence sym a -> Doc ann
C.prettySymSequence (IntrinsicPrinters sym
-> TypeRepr tp1 -> RegValue' sym tp1 -> Doc ann
forall sym (tp :: CrucibleType) ann.
IsExpr (SymExpr sym) =>
IntrinsicPrinters sym -> TypeRepr tp -> RegValue' sym tp -> Doc ann
ppRegVal IntrinsicPrinters sym
iFns TypeRepr tp1
tp' (RegValue' sym tp1 -> Doc ann)
-> (RegValue sym tp1 -> RegValue' sym tp1)
-> RegValue sym tp1
-> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall sym (tp :: CrucibleType).
RegValue sym tp -> RegValue' sym tp
C.RV @sym) SymSequence sym (RegValue sym tp1)
RegValue sym tp
v
    C.StructRepr CtxRepr ctx
fields ->
      let struct :: [Doc ann] -> Doc ann
struct = Doc ann -> Doc ann -> Doc ann -> [Doc ann] -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann -> [Doc ann] -> Doc ann
PP.encloseSep Doc ann
forall ann. Doc ann
PP.lbrace Doc ann
forall ann. Doc ann
PP.rbrace (Doc ann
forall ann. Doc ann
PP.comma Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
PP.<> Doc ann
forall ann. Doc ann
PP.space) in
      let entries :: Assignment (RegEntry sym) ctx
entries = (forall (x :: CrucibleType).
 TypeRepr x -> RegValue' sym x -> RegEntry sym x)
-> CtxRepr ctx
-> Assignment (RegValue' sym) ctx
-> Assignment (RegEntry sym) ctx
forall {k} (f :: k -> *) (g :: k -> *) (h :: k -> *) (a :: Ctx k).
(forall (x :: k). f x -> g x -> h x)
-> Assignment f a -> Assignment g a -> Assignment h a
Ctx.zipWith (\TypeRepr x
t (C.RV RegValue sym x
fv) -> forall sym (tp :: CrucibleType).
TypeRepr tp -> RegValue sym tp -> RegEntry sym tp
C.RegEntry @sym TypeRepr x
t RegValue sym x
fv) CtxRepr ctx
fields RegValue sym tp
Assignment (RegValue' sym) ctx
v in
      [Doc ann] -> Doc ann
forall {ann}. [Doc ann] -> Doc ann
struct ((forall (x :: CrucibleType). RegEntry sym x -> Doc ann)
-> forall (x :: Ctx CrucibleType).
   Assignment (RegEntry sym) x -> [Doc ann]
forall k l (t :: (k -> *) -> l -> *) (f :: k -> *) a.
FoldableFC t =>
(forall (x :: k). f x -> a) -> forall (x :: l). t f x -> [a]
forall (f :: CrucibleType -> *) a.
(forall (x :: CrucibleType). f x -> a)
-> forall (x :: Ctx CrucibleType). Assignment f x -> [a]
toListFC (\(C.RegEntry TypeRepr x
tp' RegValue sym x
val) -> IntrinsicPrinters sym -> TypeRepr x -> RegValue' sym x -> Doc ann
forall sym (tp :: CrucibleType) ann.
IsExpr (SymExpr sym) =>
IntrinsicPrinters sym -> TypeRepr tp -> RegValue' sym tp -> Doc ann
ppRegVal IntrinsicPrinters sym
iFns TypeRepr x
tp' (forall sym (tp :: CrucibleType).
RegValue sym tp -> RegValue' sym tp
C.RV @sym RegValue sym x
val)) Assignment (RegEntry sym) ctx
entries)
    C.IntrinsicRepr SymbolRepr nm
nm CtxRepr ctx
tyCtx ->
      case SymbolRepr nm
-> MapF SymbolRepr (IntrinsicPrettyFn sym)
-> Maybe (IntrinsicPrettyFn sym nm)
forall {v} (k :: v -> *) (tp :: v) (a :: v -> *).
OrdF k =>
k tp -> MapF k a -> Maybe (a tp)
MapF.lookup SymbolRepr nm
nm (IntrinsicPrinters sym -> MapF SymbolRepr (IntrinsicPrettyFn sym)
forall sym.
IntrinsicPrinters sym -> MapF SymbolRepr (IntrinsicPrettyFn sym)
getIntrinsicPrinters IntrinsicPrinters sym
iFns) of
        Maybe (IntrinsicPrettyFn sym nm)
Nothing ->
          let strNm :: [Char]
strNm = Text -> [Char]
Text.unpack (SymbolRepr nm -> Text
forall (nm :: Symbol). SymbolRepr nm -> Text
C.symbolRepr SymbolRepr nm
nm) in
          -- TODO: Replace with `panic`, exception, or `Maybe`
          [Char] -> Doc ann
forall a. HasCallStack => [Char] -> a
error ([Char]
"ppRegVal: Missing pretty-printing function for intrinsic: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
List.++ [Char]
strNm)
        Just (IntrinsicPrettyFn forall (ctx :: Ctx CrucibleType) ann.
Assignment TypeRepr ctx -> Intrinsic sym nm ctx -> Doc ann
f) -> CtxRepr ctx -> Intrinsic sym nm ctx -> Doc ann
forall (ctx :: Ctx CrucibleType) ann.
Assignment TypeRepr ctx -> Intrinsic sym nm ctx -> Doc ann
f CtxRepr ctx
tyCtx Intrinsic sym nm ctx
RegValue sym tp
v

    -- More complex cases
    C.FunctionHandleRepr CtxRepr ctx
args TypeRepr ret
ret -> CtxRepr ctx -> TypeRepr ret -> FnVal sym ctx ret -> Doc ann
forall (ctx :: Ctx CrucibleType) (ret :: CrucibleType) sym ann.
CtxRepr ctx -> TypeRepr ret -> FnVal sym ctx ret -> Doc ann
ppFnVal CtxRepr ctx
args TypeRepr ret
ret FnVal sym ctx ret
RegValue sym tp
v
    TypeRepr tp
_ -> Doc ann
"<unsupported>"

ppFnVal ::
  C.CtxRepr ctx ->
  C.TypeRepr ret ->
  C.FnVal sym ctx ret ->
  PP.Doc ann
ppFnVal :: forall (ctx :: Ctx CrucibleType) (ret :: CrucibleType) sym ann.
CtxRepr ctx -> TypeRepr ret -> FnVal sym ctx ret -> Doc ann
ppFnVal CtxRepr ctx
args TypeRepr ret
ret =
  \case
    C.HandleFnVal FnHandle ctx ret
hdl ->
      [Doc ann] -> Doc ann
forall {ann}. [Doc ann] -> Doc ann
PP.hsep
      [ FunctionName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. FunctionName -> Doc ann
PP.pretty (FnHandle ctx ret -> FunctionName
forall (args :: Ctx CrucibleType) (ret :: CrucibleType).
FnHandle args ret -> FunctionName
C.handleName FnHandle ctx ret
hdl)
      , Doc ann
":"
      , [Doc ann] -> Doc ann
forall {ann}. [Doc ann] -> Doc ann
PP.hsep (Doc ann -> [Doc ann] -> [Doc ann]
forall ann. Doc ann -> [Doc ann] -> [Doc ann]
PP.punctuate Doc ann
" ->" ((forall (x :: CrucibleType). TypeRepr x -> Doc ann)
-> forall (x :: Ctx CrucibleType).
   Assignment TypeRepr x -> [Doc ann]
forall k l (t :: (k -> *) -> l -> *) (f :: k -> *) a.
FoldableFC t =>
(forall (x :: k). f x -> a) -> forall (x :: l). t f x -> [a]
forall (f :: CrucibleType -> *) a.
(forall (x :: CrucibleType). f x -> a)
-> forall (x :: Ctx CrucibleType). Assignment f x -> [a]
toListFC TypeRepr x -> Doc ann
forall ann. TypeRepr x -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall (x :: CrucibleType). TypeRepr x -> Doc ann
PP.pretty CtxRepr ctx
args))
      , Doc ann
"->"
      , TypeRepr ret -> Doc ann
forall ann. TypeRepr ret -> Doc ann
forall a ann. Pretty a => a -> Doc ann
PP.pretty TypeRepr ret
ret
      ]
    FnVal sym ctx ret
_ -> Doc ann
"<unsupported>"