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

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

module Lang.Crucible.Debug.Trace
  ( TraceEntry(..)
  , Trace
  , empty
  , append
  , latest
  ) where

import Control.Lens qualified as Lens
import Data.Parameterized.Classes (ixF')
import Data.Parameterized.Context qualified as Ctx
import Data.Parameterized.NatRepr qualified as NatRepr
import Data.Parameterized.NatRepr (type (<=), NatRepr)
import Data.Parameterized.Some (Some(Some))
import Data.RingBuffer qualified as RB
import Data.Vector qualified as V
import Lang.Crucible.CFG.Core qualified as C
import Lang.Crucible.CFG.Extension qualified as C
import Prettyprinter qualified as PP
import What4.ProgramLoc qualified as W4

data TraceEntry ext
  = forall blocks init ret.
    TraceEntry
    { ()
traceCfg :: C.CFG ext blocks init ret
    , ()
traceBlock :: Some (C.BlockID blocks)
    }

instance C.PrettyExt ext => PP.Pretty (TraceEntry ext) where
  pretty :: forall ann. TraceEntry ext -> Doc ann
pretty (TraceEntry CFG ext blocks init ret
cfg (Some BlockID blocks x
blkId)) =
    let blk :: Block ext blocks ret x
blk = CFG ext blocks init ret -> BlockMap ext blocks ret
forall ext (blocks :: Ctx (Ctx CrucibleType))
       (init :: Ctx CrucibleType) (ret :: CrucibleType).
CFG ext blocks init ret -> BlockMap ext blocks ret
C.cfgBlockMap CFG ext blocks init ret
cfg BlockMap ext blocks ret
-> Getting
     (Block ext blocks ret x)
     (BlockMap ext blocks ret)
     (Block ext blocks ret x)
-> Block ext blocks ret x
forall s a. s -> Getting a s a -> a
Lens.^. IndexF (BlockMap ext blocks ret) x
-> Lens'
     (BlockMap ext blocks ret) (IxValueF (BlockMap ext blocks ret) x)
forall k m (x :: k).
IxedF' k m =>
IndexF m x -> Lens' m (IxValueF m x)
forall (x :: Ctx CrucibleType).
IndexF (BlockMap ext blocks ret) x
-> Lens'
     (BlockMap ext blocks ret) (IxValueF (BlockMap ext blocks ret) x)
ixF' (BlockID blocks x -> Index blocks x
forall (blocks :: Ctx (Ctx CrucibleType)) (tp :: Ctx CrucibleType).
BlockID blocks tp -> Index blocks tp
C.blockIDIndex BlockID blocks x
blkId) in
    Bool -> Size x -> StmtSeq ext blocks ret x -> Doc ann
forall ext (ctx :: Ctx CrucibleType)
       (blocks :: Ctx (Ctx CrucibleType)) (ret :: CrucibleType) ann.
PrettyExt ext =>
Bool -> Size ctx -> StmtSeq ext blocks ret ctx -> Doc ann
ppStmtSeq Bool
True (Assignment TypeRepr x -> Size x
forall {k} (f :: k -> *) (ctx :: Ctx k).
Assignment f ctx -> Size ctx
Ctx.size (Block ext blocks ret x -> Assignment TypeRepr x
forall ext (blocks :: Ctx (Ctx CrucibleType)) (ret :: CrucibleType)
       (ctx :: Ctx CrucibleType).
Block ext blocks ret ctx -> CtxRepr ctx
C.blockInputs Block ext blocks ret x
blk)) (Block ext blocks ret x -> StmtSeq ext blocks ret x
forall ext (blocks :: Ctx (Ctx CrucibleType)) (ret :: CrucibleType)
       (ctx :: Ctx CrucibleType).
Block ext blocks ret ctx -> StmtSeq ext blocks ret ctx
C._blockStmts Block ext blocks ret x
blk)
    -- TODO: Taken from Crucible. Export upstream?
    where
      prefixLineNum :: Bool -> W4.ProgramLoc -> PP.Doc ann -> PP.Doc ann
      prefixLineNum :: forall ann. Bool -> ProgramLoc -> Doc ann -> Doc ann
prefixLineNum Bool
True ProgramLoc
pl Doc ann
d = [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
PP.vcat [String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
PP.pretty String
"%" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
PP.<+> Position -> Doc ann
forall ann. Position -> Doc ann
W4.ppNoFileName (ProgramLoc -> Position
W4.plSourceLoc ProgramLoc
pl), Doc ann
d]
      prefixLineNum Bool
False ProgramLoc
_ Doc ann
d = Doc ann
d

      ppStmtSeq :: C.PrettyExt ext => Bool -> Ctx.Size ctx -> C.StmtSeq ext blocks ret ctx -> PP.Doc ann
      ppStmtSeq :: forall ext (ctx :: Ctx CrucibleType)
       (blocks :: Ctx (Ctx CrucibleType)) (ret :: CrucibleType) ann.
PrettyExt ext =>
Bool -> Size ctx -> StmtSeq ext blocks ret ctx -> Doc ann
ppStmtSeq Bool
ppLineNum Size ctx
h (C.ConsStmt ProgramLoc
pl Stmt ext ctx ctx'
s StmtSeq ext blocks ret ctx'
r) =
        [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
PP.vcat
        [ Bool -> ProgramLoc -> Doc ann -> Doc ann
forall ann. Bool -> ProgramLoc -> Doc ann -> Doc ann
prefixLineNum Bool
ppLineNum ProgramLoc
pl (Size ctx -> Stmt ext ctx ctx' -> Doc ann
forall ext (ctx :: Ctx CrucibleType) (ctx' :: Ctx CrucibleType)
       ann.
PrettyExt ext =>
Size ctx -> Stmt ext ctx ctx' -> Doc ann
C.ppStmt Size ctx
h Stmt ext ctx ctx'
s)
        , Bool -> Size ctx' -> StmtSeq ext blocks ret ctx' -> Doc ann
forall ext (ctx :: Ctx CrucibleType)
       (blocks :: Ctx (Ctx CrucibleType)) (ret :: CrucibleType) ann.
PrettyExt ext =>
Bool -> Size ctx -> StmtSeq ext blocks ret ctx -> Doc ann
ppStmtSeq Bool
ppLineNum (Size ctx -> Stmt ext ctx ctx' -> Size ctx'
forall (ctx :: Ctx CrucibleType) ext (ctx' :: Ctx CrucibleType).
Size ctx -> Stmt ext ctx ctx' -> Size ctx'
C.nextStmtHeight Size ctx
h Stmt ext ctx ctx'
s) StmtSeq ext blocks ret ctx'
r
        ]
      ppStmtSeq Bool
ppLineNum Size ctx
_ (C.TermStmt ProgramLoc
pl TermStmt blocks ret ctx
s) =
        Bool -> ProgramLoc -> Doc ann -> Doc ann
forall ann. Bool -> ProgramLoc -> Doc ann -> Doc ann
prefixLineNum Bool
ppLineNum ProgramLoc
pl (TermStmt blocks ret ctx -> Doc ann
forall ann. TermStmt blocks ret ctx -> Doc ann
forall a ann. Pretty a => a -> Doc ann
PP.pretty TermStmt blocks ret ctx
s)


newtype Trace ext = Trace (RB.RingBuffer V.Vector (TraceEntry ext))

-- 'RB.new' requires a nonzero input, which we enforce with types
empty :: (1 <= n) => NatRepr n -> IO (Trace ext)
empty :: forall (n :: Natural) ext. (1 <= n) => NatRepr n -> IO (Trace ext)
empty NatRepr n
n = RingBuffer Vector (TraceEntry ext) -> Trace ext
forall ext. RingBuffer Vector (TraceEntry ext) -> Trace ext
Trace (RingBuffer Vector (TraceEntry ext) -> Trace ext)
-> IO (RingBuffer Vector (TraceEntry ext)) -> IO (Trace ext)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO (RingBuffer Vector (TraceEntry ext))
forall (v :: * -> *) a. Vector v a => Int -> IO (RingBuffer v a)
RB.new (Natural -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (NatRepr n -> Natural
forall (n :: Natural). NatRepr n -> Natural
NatRepr.natValue NatRepr n
n))

-- | O(1).
append :: Trace ext -> TraceEntry ext -> IO ()
append :: forall ext. Trace ext -> TraceEntry ext -> IO ()
append (Trace RingBuffer Vector (TraceEntry ext)
b) TraceEntry ext
e = TraceEntry ext -> RingBuffer Vector (TraceEntry ext) -> IO ()
forall (v :: * -> *) a. Vector v a => a -> RingBuffer v a -> IO ()
RB.append TraceEntry ext
e RingBuffer Vector (TraceEntry ext)
b
{-# INLINEABLE append #-}

-- | Get up to the @N@ latest entries
latest :: Trace ext -> Int -> IO [TraceEntry ext]
latest :: forall ext. Trace ext -> Int -> IO [TraceEntry ext]
latest (Trace RingBuffer Vector (TraceEntry ext)
b) Int
n = Int -> IO [TraceEntry ext]
go Int
n
  where
    go :: Int -> IO [TraceEntry ext]
go Int
x | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = [TraceEntry ext] -> IO [TraceEntry ext]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    go Int
x = do
      Maybe (TraceEntry ext)
mEnt <- RingBuffer Vector (TraceEntry ext)
-> Int -> IO (Maybe (TraceEntry ext))
forall (v :: * -> *) a.
Vector v a =>
RingBuffer v a -> Int -> IO (Maybe a)
RB.latest RingBuffer Vector (TraceEntry ext)
b (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
x)
      case Maybe (TraceEntry ext)
mEnt of
        Maybe (TraceEntry ext)
Nothing -> [TraceEntry ext] -> IO [TraceEntry ext]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
        Just TraceEntry ext
ent -> (TraceEntry ext
ent TraceEntry ext -> [TraceEntry ext] -> [TraceEntry ext]
forall a. a -> [a] -> [a]
:) ([TraceEntry ext] -> [TraceEntry ext])
-> IO [TraceEntry ext] -> IO [TraceEntry ext]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO [TraceEntry ext]
go (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)