Halt gracefully

app/Main.hs

@@ -22,7 +22,7 @@ import System.Exit (exitWith, ExitCode(..))
 import Types (Word)
 import Util
 import System.IO
-import Data.Maybe (fromMaybe, isJust)
+import Data.Maybe (fromMaybe, isJust, isNothing)
 import qualified Prelude as P
 import Prelude hiding (Ordering (..), Word, break, init, log, map, not, repeat, undefined, (&&), (++), (||), replicate, zip, take)
 import Data.Traversable
@@ -241,10 +241,10 @@ runNormalMemory Options{..} elf entryOffset leakOutputHandle leakDigest finalSta
           (s', o) <- step i s
           let halted = Core.stateHalt s'
           (mi', sysExit) <- case halted of
-            Core.Running -> do
+            Nothing -> do
               mi'' <- next s' o
               pure (mi'', False)
-            Core.Syscall -> do
+            Just (Core.Syscall resumePc) -> do
               -- Handle syscall, write return value to a0, resume
               let serialize = if idx == 0
                     then BS.toStrict . encode
@@ -254,7 +254,9 @@ runNormalMemory Options{..} elf entryOffset leakOutputHandle leakDigest finalSta
               case mRet of
                 Nothing -> pure (Nothing, True)  -- exit
                 Just ret -> do
-                  let s'' = s' {Core.stateHalt = Core.Running,
+                  let s'' = s' {Core.stateHalt = Nothing,
+                                Core.stateFePc = resumePc,
+                                Core.stateDePc = resumePc,
                                 Core.stateRegFile = modifyRF 10 ret (Core.stateRegFile s')}
                   mi'' <- next s'' o
                   pure (mi'', False)
@@ -270,7 +272,7 @@ runNormalMemory Options{..} elf entryOffset leakOutputHandle leakDigest finalSta
           { stepMem = mem
           , stepSim = newSim
           , stepNextInput = mi'
-          , stepContinue = not sysExit && (halted == Core.Running || halted == Core.Syscall)
+          , stepContinue = not sysExit && (isNothing halted || (\case { Just (Core.Syscall _) -> True; _ -> False }) halted)
           , stepFinalState = s'
           , stepLeakage = leakOutput
           }
@@ -336,7 +338,7 @@ runExecutable opts@Options{..} = do
                 }
           finalState <- readIORef finalStateRef
           case finalState of
-            Just state | Core.stateHalt state == Core.SecurityViolation -> do
+            Just state | Core.stateHalt state == Just Core.SecurityViolation -> do
               putStrLn "Program aborted due to security violation"
             _ -> pure ()
         else do

src/Core.hs

@@ -1,6 +1,7 @@
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE StandaloneDeriving #-}
 {-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE DerivingStrategies #-}
 
 module Core
   ( initInput,
@@ -80,7 +81,7 @@
 deriving instance (Generic (f Word), NFDataX (f Word)) => NFDataX (MemAccess f)
 
 -- | The output of the CPU.
-data Output f = Output
+newtype Output f = Output
   { -- | A memory access.
     outMem :: First (MemAccess f)
   }
@@ -89,7 +90,7 @@
 
 deriving instance Generic (Output f)
 
 deriving instance (Generic (f Word), NFDataX (f Word)) => NFDataX (Output f)
 
 instance Semigroup (Output f) where
   Output mem <> Output mem' =
@@ -99,7 +100,7 @@
   mempty = Output mempty
 
 -- | CPU halt state
-data HaltState = Running | EBreak | Syscall | SecurityViolation
+data HaltState = EBreak Address | Syscall Address | SecurityViolation
   deriving (Show, Eq, Generic)
 
 instance NFDataX HaltState
@@ -129,7 +130,9 @@
     -- | Control/forwarding lines.
     stateCtrl :: Control f,
     -- | CPU halt state
-    stateHalt :: HaltState
+    stateHalt :: Maybe HaltState,
+    -- | Pending halt state (propagating through pipeline to ensure flush)
+    stateHaltPending :: Maybe HaltState
   }
 
 deriving instance (Show (f Word)) => Show (State f)
@@ -175,7 +178,7 @@
 -- | No secrets here, buddy: unwrap a word. If it's public, we gucci. If it's
 -- private, die.
 noSecrets' :: (Access f) => f a -> b -> (a -> CPUM f b) -> CPUM f b
-noSecrets' w a m = noSecrets w (setSecurityViolation >> pure a) m
+noSecrets' w a = noSecrets w (setSecurityViolation >> pure a)
 
 -- | Run the CPU for one step.
 circuit :: (Access f) => State f -> Input f -> (State f, Output f)
@@ -211,7 +214,8 @@
       stateWbAluRes = pure 0,
       stateRegFile = initRF,
       stateCtrl = initCtrl,
-      stateHalt = Running
+      stateHalt = Nothing,
+      stateHaltPending = Nothing
     }
 
 -- | Initial control lines.
@@ -231,36 +235,29 @@
 withCtrlReset m = do
   modify $ \s -> s {stateCtrl = initCtrl}
   m
-  ctrl <- gets stateCtrl
-  pure ctrl
-
--- | Stop the CPU due to ebreak.
-halt :: CPUM f ()
-halt =
-  modify $ \s -> s {stateHalt = EBreak}
-
--- | Stop the CPU due to syscall.
-setSyscall :: CPUM f ()
-setSyscall =
-  modify $ \s -> s {stateHalt = Syscall}
+  gets stateCtrl
 
 -- | Set security violation flag.
 setSecurityViolation :: CPUM f ()
 setSecurityViolation =
-  modify $ \s -> s {stateHalt = SecurityViolation}
+  modify $ \s -> s {stateHalt = Just SecurityViolation}
 
 -- | The fetch stage.
 fetch :: CPUM f ()
 fetch = do
   pc <- gets stateFePc
   ctrl <- gets stateCtrl
+  status <- gets stateHalt
+  pending <- gets stateHaltPending
+
+  let isHalted = isJust status || isJust pending
 
-  -- Always try to read unless the instruction in the `memory` stage is a load or a store.
-  unless (ctrlMeMemInstr ctrl) $
+  -- Always try to read unless the instruction in the `memory` stage is a load or a store, or the core is halting/halted.
+  unless (ctrlMeMemInstr ctrl || isHalted) $
     readPC pc
-  
-  -- We stall if the instruction in the `memory` stage is a load or a store.
-  let stall = ctrlMeMemInstr ctrl
+
+  -- We stall if the instruction in the `memory` stage is a load or a store, or the core is halting/halted.
+  let stall = ctrlMeMemInstr ctrl || isHalted
 
   let next_pc =
         fromMaybe
@@ -274,7 +271,7 @@
     s -- Increment program counter for next fetch.
       { stateFePc = next_pc,
         -- Propagate program counter to next stage.
-        stateDePc = pc
+        stateDePc = if stall then stateDePc s else pc
       }
 
 -- | Decode stage.
@@ -283,9 +280,10 @@
   input <- ask
   ctrl <- gets stateCtrl
   status <- gets stateHalt
+  pending <- gets stateHaltPending
 
   ir <-
-    if (inputIsInstr input)
+    if inputIsInstr input
       then noSecrets' (inputMem input) (Nop Halted) (pure . decode')
       else pure $ Nop MemoryBusBusy
 
@@ -309,7 +307,7 @@
         -- If a break is executed in this cycle, we halt.
         | break_current_cycle = Nop Halted
         -- If the core is not running anymore, we halt.
-        | status /= Running = Nop Halted
+        | isJust status || isJust pending = Nop Halted
         -- Otherwise we process the decoded instruction.
         | otherwise = ir
 
@@ -395,12 +393,18 @@
             PC -> gets $ pack . stateExPc
           let imm' = imm ++# (0 :: BitVector 12)
           pure (ADD, pure base', pure imm')
-        Instruction.IType (Env Call) _ _ _ ->
-          lift halt >> empty
-        Instruction.IType (Env Break) _ _ _ ->
-          lift halt >> empty
+        Instruction.IType (Env Call) _ _ _ -> do
+          pc <- gets stateExPc
+          pendingHalt (Syscall (pc + 4))
+        Instruction.IType (Env Break) _ _ _ -> do
+          pc <- gets stateExPc
+          pendingHalt (EBreak (pc + 4))
         Instruction.Nop _ -> empty
 
+    pendingHalt :: HaltState -> MaybeT (CPUM f) a
+    pendingHalt hState = do
+      lift (modify $ \s -> s {stateHaltPending = Just hState}) >> empty
+
     rs1 :: MaybeT (CPUM f) (f Word)
     rs1 = do
       ir <- gets stateExInstr
@@ -464,6 +468,11 @@
   ir <- gets stateMeInstr
   res <- gets stateMeAluRes
   val <- gets stateMeStoreRes
+  pending <- gets stateHaltPending
+
+  case pending of
+    Just hlt -> modify $ \s -> s {stateHalt = Just hlt, stateHaltPending = Nothing}
+    Nothing -> pure ()
 
   modify $ \s ->
     s {stateWbInstr = ir, stateWbAluRes = res}
@@ -492,12 +501,6 @@
       setLines $ \c -> c {ctrlMeRegFwd = Just (rd, res)}
     Instruction.UType _ rd _ ->
       setLines $ \c -> c {ctrlMeRegFwd = Just (rd, res)}
-    Instruction.IType (Env Call) _ _ _ -> do
-      setSyscall
-    Instruction.IType (Env Break) _ _ _ -> do
-      halt
-    Instruction.Nop Halted -> do
-      halt    
     _ -> pure ()
 
 -- | Commit computations to the register file.
@@ -506,7 +509,6 @@
   input <- asks inputMem
   ir <- gets stateWbInstr
   res <- gets stateWbAluRes
-  status <- gets stateHalt
 
   case ir of
     Instruction.RType _ rd _ _ -> do

src/Elf/ElfLoader.hs

@@ -21,14 +21,14 @@ import Data.Char (chr)
 import Data.Elf
 import Data.Elf.Constants
 import Data.Elf.Headers
-import Data.Int
 import Data.Word
 import Types
 import RegFile (modifyRF)
 import Memory.Types
 import Util
 import Prelude hiding (Ordering (..), Word, init, log, map, not, repeat, undefined, (&&), (++), (||))
 import qualified Prelude as P
+import Control.Monad.IO.Class (MonadIO)
 
 loadElf :: (MonadMemory m) => Elf -> m ()
 loadElf elf@(Elf classS _) = withSingElfClassI classS $ do
@@ -96,27 +96,29 @@ readStringFromMemory addr count = do
 -- | Called when the core halts with a syscall. Return the value to write to a0, or Nothing to truly exit.
 type Instrument f m = Core.State f -> m (Maybe (f Types.Word))
 
-runElf :: forall f m. (Access f, MonadMemory m) => Instrument f m -> CircuitSim m (Core.Input f) (Core.State f) (Core.Output f) -> m ()
+runElf :: forall f m. (Access f, MonadIO m, MonadMemory m) => Instrument f m -> CircuitSim m (Core.Input f) (Core.State f) (Core.Output f) -> m ()
 runElf instr c = go c
   where
     go sim@(CircuitSim i s step next) = do
       (s', o) <- step i s
       case Core.stateHalt s' of
-        Core.Running -> do
+        Nothing -> do
           mi' <- next s' o
           case mi' of
             Just i' -> go $ sim {circuitInput = i', circuitState = s'}
             Nothing -> pure ()
-        Core.Syscall -> do
+        Just (Core.Syscall resumePc) -> do
           mRet <- instr s'
           case mRet of
             Nothing -> pure ()
             Just ret -> do
-              let s'' = s' {Core.stateHalt = Core.Running,
+              let s'' = s' {Core.stateHalt = Nothing,
+                            Core.stateFePc = resumePc,
+                            Core.stateDePc = resumePc,
                             Core.stateRegFile = modifyRF 10 ret (Core.stateRegFile s')}
               mi' <- next s'' o
               case mi' of
                 Just i' -> go $ sim {circuitInput = i', circuitState = s''}
                 Nothing -> pure ()
-        Core.EBreak -> pure ()
-        Core.SecurityViolation -> pure ()
+        Just (Core.EBreak _) -> pure ()
+        Just Core.SecurityViolation -> pure ()

src/Leak/MonitorPC/PC.hs

@@ -17,7 +17,7 @@ import Control.Monad
 import Core (Input (..), MemAccess (..), Output (..))
 import qualified Core
 import Data.Functor.Identity
-import Data.Maybe (isJust)
+import Data.Maybe (isJust, isNothing)
 import Data.Monoid
 import qualified Leak.MonitorPC.MonitorLeak as Leak
 import qualified Leak.MonitorPC.Sim as Sim
@@ -43,7 +43,7 @@ proj :: Core.State Identity -> (((), Core.State Identity), Sim.State)
 proj s = (ts, ss)
   where
     ts = Leak.leakProject Leak.monitorPC s
-    halted = Core.stateHalt s /= Core.Running
+    halted = isJust (Core.stateHalt s)
     ss =
       Sim.State
         { Sim.stateFePc = if halted then 0 else Core.stateFePc s,
@@ -58,7 +58,7 @@ proj s = (ts, ss)
           Sim.stateStallFetch = not halted && toStallFetch (Core.stateCtrl s),
           Sim.stateStallDecode = not halted && toStallDecode (Core.stateCtrl s),
           Sim.stateJumpAddr = if halted then Nothing else Core.ctrlExAddress $ Core.stateCtrl s,
-          Sim.stateFirstCycle = not halted && Core.stateHalt s == Core.Running
+          Sim.stateFirstCycle = not halted && isNothing (Core.stateHalt s)
         }
 
     killJump :: Leak.Instr -> Leak.Instr