TypedTrace.hs

{-# LANGUAGE GADTs               #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- |
-- Copyright : (c) NASA, 2024-2025
-- License   : BSD-style (see the LICENSE file in the distribution)
--
-- Read and modify traces of specifications.
--
-- In the context of this module, a trace refers only to the values of
-- externs, since the state of a spec can be reproduced if the values of
-- externs are known.
module Copilot.Visualize.TypedTrace
    ( extractTrace
    , updateWithTrace
    , Trace(..)
    , UValues(..)
    )
  where

-- External imports
import           Data.List          hiding ((++))
import           Data.Maybe         (fromMaybe)
import qualified Data.Type.Equality as DE
import           Data.Typeable      (Typeable)
import           Prelude            hiding (div, not, (++), (<), (>))

-- External imports: Copilot
import qualified Copilot.Core     as Core
import           Language.Copilot hiding (interpret, typeOf)

-- | Map stream names to their values.
newtype Trace = Trace
  { traceMap :: [ (String, UValues) ]
  }

-- | Existentially typed list of values of a stream.
data UValues = forall a . Typeable a => UValues
  { uvType   :: Core.Type a
  , uvValues :: [ a ]
  }

-- | Extract a trace from a core specification.
extractTrace :: Core.Spec -> Trace
extractTrace spec = Trace $ concat $ concat
  [ fmap extractTraceStream (Core.specStreams spec)
  , fmap extractTraceObserver (Core.specObservers spec)
  , fmap extractTraceTrigger (Core.specTriggers spec)
  ]

-- | Extract the values of a core 'Stream'.
extractTraceStream :: Core.Stream -> [ (String, UValues) ]
extractTraceStream (Core.Stream _id _buf expr _ty) =
  extractTraceExpr expr

-- | Extract the values of a core 'Observer'.
extractTraceObserver :: Core.Observer -> [ (String, UValues) ]
extractTraceObserver (Core.Observer _name expr _ty) =
  extractTraceExpr expr

-- | Extract the values of a core 'Trigger'.
extractTraceTrigger :: Core.Trigger -> [ (String, UValues) ]
extractTraceTrigger (Core.Trigger _name expr args) = concat $
    extractTraceExpr expr
  : fmap extractTraceUExpr args

-- | Extract the values of a core 'Expr'.
extractTraceExpr :: Core.Expr a -> [ (String, UValues) ]
extractTraceExpr (Core.Local _ _ _ expr1 expr2) = concat
  [ extractTraceExpr expr1
  , extractTraceExpr expr2
  ]
extractTraceExpr (Core.ExternVar ty name values) =
  [ (name, UValues ty (fromMaybe [] values)) ]
extractTraceExpr (Core.Op1 _op expr) =
  extractTraceExpr expr
extractTraceExpr (Core.Op2 _op expr1 expr2) = concat
  [ extractTraceExpr expr1
  , extractTraceExpr expr2
  ]
extractTraceExpr (Core.Op3 _op expr1 expr2 expr3) = concat
  [ extractTraceExpr expr1
  , extractTraceExpr expr2
  , extractTraceExpr expr3
  ]
extractTraceExpr (Core.Label _ty _lbl expr) =
  extractTraceExpr expr
extractTraceExpr _ = []

-- | Extract the values of a core 'UExpr'.
extractTraceUExpr :: Core.UExpr -> [ (String, UValues) ]
extractTraceUExpr (Core.UExpr ty expr) = extractTraceExpr expr

-- | Update externs in a spec.
updateWithTrace :: Trace -> Core.Spec -> Core.Spec
updateWithTrace trace spec = spec
  { Core.specStreams =
      fmap (updateWithTraceStream trace) (Core.specStreams spec)
  , Core.specObservers =
      fmap (updateWithTraceObserver trace) (Core.specObservers spec)
  , Core.specTriggers =
      fmap (updateWithTraceTrigger trace) (Core.specTriggers spec)
  }

-- | Update the values of a 'Stream' based on an input trace.
updateWithTraceStream :: Trace -> Core.Stream -> Core.Stream
updateWithTraceStream trace (Core.Stream ident buf expr ty) =
  Core.Stream ident buf (updateWithTraceExpr trace expr) ty

-- | Update the values of an 'Observer' based on an input trace.
updateWithTraceObserver :: Trace -> Core.Observer -> Core.Observer
updateWithTraceObserver trace (Core.Observer name expr ty) =
  Core.Observer name (updateWithTraceExpr trace expr) ty

-- | Update the values of a 'Trigger' based on an input trace.
updateWithTraceTrigger :: Trace -> Core.Trigger -> Core.Trigger
updateWithTraceTrigger trace (Core.Trigger name expr args) =
  Core.Trigger
    name
    (updateWithTraceExpr trace expr)
    (fmap (updateWithTraceUExpr trace) args)

-- | Update the values of an 'Expr' based on an input trace.
updateWithTraceExpr :: Trace -> Core.Expr a -> Core.Expr a
updateWithTraceExpr trace (Core.Local ty1 ty2 name expr1 expr2) =
  Core.Local
    ty1
    ty2
    name
    (updateWithTraceExpr trace expr1)
    (updateWithTraceExpr trace expr2)
updateWithTraceExpr trace (Core.ExternVar ty name values) =
    Core.ExternVar ty name values'
  where
    values' | Just (UValues ty2 vals) <- lookup name (traceMap trace)
            , Just DE.Refl <- DE.testEquality ty ty2
            = Just vals
            | otherwise
            = values
updateWithTraceExpr trace (Core.Op1 op expr) =
  Core.Op1 op (updateWithTraceExpr trace expr)
updateWithTraceExpr trace (Core.Op2 op expr1 expr2) =
  Core.Op2 op
    (updateWithTraceExpr trace expr1)
    (updateWithTraceExpr trace expr2)
updateWithTraceExpr trace (Core.Op3 op expr1 expr2 expr3) =
  Core.Op3 op
    (updateWithTraceExpr trace expr1)
    (updateWithTraceExpr trace expr2)
    (updateWithTraceExpr trace expr3)
updateWithTraceExpr trace (Core.Label ty lbl expr) =
  Core.Label ty lbl (updateWithTraceExpr trace expr)
updateWithTraceExpr trace x = x

-- | Update the values of a 'UExpr' based on an input trace.
updateWithTraceUExpr :: Trace -> Core.UExpr -> Core.UExpr
updateWithTraceUExpr trace (Core.UExpr ty expr) =
  Core.UExpr ty (updateWithTraceExpr trace expr)