diff --git a/spec/haskell/src/SEL4/Kernel/Thread/AARCH64.hs b/spec/haskell/src/SEL4/Kernel/Thread/AARCH64.hs
index 3aad6af2ec..1b6ce2996d 100644
--- a/spec/haskell/src/SEL4/Kernel/Thread/AARCH64.hs
+++ b/spec/haskell/src/SEL4/Kernel/Thread/AARCH64.hs
@@ -19,7 +19,6 @@ import {-# SOURCE #-} SEL4.Kernel.Init
 import {-# SOURCE #-} SEL4.Object.TCB (asUser, threadGet)
 import SEL4.Machine.RegisterSet.AARCH64(Register(..))
 import SEL4.Object.VCPU.AARCH64
-import SEL4.Object.Structures.AARCH64
 import SEL4.Object.FPU.AARCH64
 
 switchToThread :: PPtr TCB -> Kernel ()
diff --git a/spec/haskell/src/SEL4/Object/FPU/AARCH64.hs b/spec/haskell/src/SEL4/Object/FPU/AARCH64.hs
index 08114ac63f..ac4a4266f1 100644
--- a/spec/haskell/src/SEL4/Object/FPU/AARCH64.hs
+++ b/spec/haskell/src/SEL4/Object/FPU/AARCH64.hs
@@ -33,10 +33,10 @@ switchLocalFpuOwner :: Maybe (PPtr TCB) -> Kernel ()
 switchLocalFpuOwner newOwner = do
     doMachineOp enableFpu
     curFpuOwner <- gets (armKSCurFPUOwner . ksArchState)
-    when (isJust curFpuOwner) (saveFpuState (fromJust curFpuOwner))
-    if isJust newOwner
-      then loadFpuState (fromJust newOwner)
-      else doMachineOp disableFpu
+    maybe (return ()) saveFpuState curFpuOwner
+    case newOwner of
+        Nothing -> doMachineOp disableFpu
+        Just tcbPtr -> loadFpuState tcbPtr
     modifyArchState (\s -> s { armKSCurFPUOwner = newOwner })
 
 fpuRelease :: PPtr TCB -> Kernel ()
diff --git a/spec/haskell/src/SEL4/Object/TCB/AARCH64.hs b/spec/haskell/src/SEL4/Object/TCB/AARCH64.hs
index 113e012784..805ab06d3d 100644
--- a/spec/haskell/src/SEL4/Object/TCB/AARCH64.hs
+++ b/spec/haskell/src/SEL4/Object/TCB/AARCH64.hs
@@ -58,6 +58,8 @@ postSetFlags :: PPtr TCB -> TcbFlags -> Kernel ()
 postSetFlags t flags =
     when (tcbFlagToWord (ArchFlag FpuDisabled) .&. flags /= 0) (fpuRelease t)
 
+-- Save and clear FPU state before setting the domain of a TCB, to ensure that
+-- we do not later write to cross-domain state.
 prepareSetDomain :: PPtr TCB -> Domain -> Kernel ()
 prepareSetDomain t newDom = do
     curDom <- curDomain