use reoptPrepareForRecovery where possible

reopt/Main_reopt.hs

@@ -12,10 +12,7 @@ import Data.ByteString qualified as BS
 import Data.ByteString.Builder qualified as Builder
 import Data.ByteString.Char8 qualified as BSC
 import Data.Either (rights)
-import Data.ElfEdit (
-  ElfSection (elfSectionAddr, elfSectionData, elfSectionName),
-  elfSections,
- )
+import Data.ElfEdit qualified as ELF
 import Data.ElfEdit qualified as Elf
 import Data.Generics.Labels ()
 import Data.IORef (IORef, modifyIORef', newIORef, readIORef)
@@ -26,7 +23,7 @@ import Data.Version (Version (versionBranch))
 import Data.Word (Word64)
 import GHC.Generics (Generic)
 import Numeric (readHex)
-import Options.Applicative
+import Options.Applicative hiding (header)
 import Prettyprinter qualified as PP
 import Prettyprinter.Render.Text qualified as PP
 import System.Exit (exitFailure)
@@ -42,18 +39,16 @@ import Text.Printf (printf)
 import Data.Macaw.Analysis.RegisterUse qualified as Macaw
 import Data.Macaw.DebugLogging qualified as Macaw
 import Data.Macaw.Discovery qualified as Macaw
-import Data.Parameterized.Some (Some (Some))
 
 import Reopt
-import Reopt.ELFArchInfo (getElfArchInfo)
 import Reopt.EncodeInvariants (
   encodeInvariantFailedMsg,
   encodeInvariantMsg,
  )
 import Reopt.Events (
   FunId (FunId),
   ReoptFunStep (AnnotationGeneration, InvariantInference),
-  ReoptGlobalStep (DiscoveryInitialization, Relinking),
+  ReoptGlobalStep (Relinking),
   ReoptLogEvent (ReoptFunStepFailed, ReoptFunStepFinished),
   joinLogEvents,
   printLogEvent,
@@ -76,7 +71,6 @@ import Reopt.Utils.Exit (
 import Reopt.VCG.Annotations as Ann (ModuleAnnotations (..))
 import Reopt.X86 (
   X86OS,
-  osArchitectureInfo,
   osLinkName,
   osPersonality,
  )
@@ -650,13 +644,13 @@ dumpDisassembly args elfPath = do
   unless (null l) $ do
     displayError "Recoverable errors occurred in reading elf file:"
     mapM_ (IO.hPrint IO.stderr) l
-  let sections = filter isCodeSection $ e ^.. elfSections
+  let sections = filter isCodeSection $ e ^.. ELF.elfSections
   when (null sections) $ do
     displayError "Binary contains no executable sections"
     exitFailure
   writeOutput (outputPath args) $ \h -> do
     forM_ sections $ \s -> do
-      printX86SectionDisassembly h (elfSectionName s) (elfSectionAddr s) (elfSectionData s)
+      printX86SectionDisassembly h (ELF.elfSectionName s) (ELF.elfSectionAddr s) (ELF.elfSectionData s)
 
 loadOptions :: Args -> LoadOptions
 loadOptions args = LoadOptions{loadOffset = loadBaseAddress args}
@@ -681,29 +675,24 @@ argsReoptOptions args = do
 showCFG :: Args -> FilePath -> IO String
 showCFG args elfPath = do
   reoptOpts <- argsReoptOptions args
-  Some hdrInfo <- do
+  hdrInfo <- do
     bs <- checkedReadFile elfPath
     case Elf.decodeElfHeaderInfo bs of
       Left (_, msg) -> do
         displayError $ "Error reading " ++ elfPath ++ ":"
         displayError $ "  " ++ msg
         exitFailure
       Right (Elf.SomeElf hdr) ->
-        pure $! Some hdr
-  let hdr = Elf.header hdrInfo
-  -- Get architecture specific information
-  marchInfo <- getElfArchInfo (Elf.headerClass hdr) (Elf.headerMachine hdr) (Elf.headerOSABI hdr)
-  (w, SomeArch ainfo pltFn) <- handleEitherStringWithExit marchInfo
-  mapM_ displayError w
+        case ELF.headerClass (ELF.header hdr) of
+          ELF.ELFCLASS32 -> do
+            displayError "Encountered 32-bit executable, reopt currently only supports X84-64"
+            exitFailure
+          ELF.ELFCLASS64 -> pure $! hdr
   mr <-
     runReoptM printLogEvent $ do
-      -- Resolve header annotations
       hdrAnn <- resolveHeader (headerPath args) (clangPath args)
-      -- Perform Discovery
-      globalStepStarted DiscoveryInitialization
-      initState <- reoptRunInit $ doInit (loadOptions args) hdrInfo ainfo pltFn reoptOpts
-      (_, discState) <- doDiscovery hdrAnn hdrInfo ainfo initState reoptOpts
-      -- Print discovery
+      (_, _, _, discState) <-
+        reoptInitialDiscovery (loadOptions args) reoptOpts hdrAnn hdrInfo
       pure $ show $ Macaw.ppDiscoveryStateBlocks discState
   handleEitherWithExit mr
 
@@ -724,13 +713,8 @@ showConstraints args elfPath = do
       funPrefix :: BSC.ByteString
       funPrefix = unnamedFunPrefix args
 
-    (os, initState) <- reoptX86Init (loadOptions args) rOpts origElf
-    let symAddrMap = initDiscSymAddrMap initState
-
-    checkSymbolUnused funPrefix symAddrMap
-
-    let ainfo = osArchitectureInfo os
-    (debugTypeMap, discState) <- doDiscovery hdrAnn origElf ainfo initState rOpts
+    (os, symAddrMap, debugTypeMap, discState) <-
+      reoptInitialDiscovery (loadOptions args) rOpts hdrAnn origElf
 
     let sysp = osPersonality os
     recoverX86Output <-
@@ -834,10 +818,7 @@ performReopt args elfPath = do
       funPrefix = unnamedFunPrefix args
 
     (os, symAddrMap, debugTypeMap, discState) <-
-      reoptPrepareForRecovery (loadOptions args) rOpts hdrAnn funPrefix origElf
-
-    when (shouldRecover args) $
-      checkSymbolUnused funPrefix symAddrMap
+      reoptInitialDiscovery (loadOptions args) rOpts hdrAnn origElf
 
     case cfgExportPath args of
       Nothing -> pure ()

src/Reopt.hs

@@ -31,7 +31,7 @@ module Reopt (
   printX86SectionDisassembly,
   doInit,
   reoptX86Init,
-  checkSymbolUnused,
+  checkNoSymbolUsesReservedPrefix,
   SomeArchitectureInfo (..),
 
   -- * Code discovery
@@ -55,7 +55,7 @@ module Reopt (
   Reopt.TypeInference.HeaderTypes.emptyAnnDeclarations,
   RecoveredModule,
   recoveredDefs,
-  reoptPrepareForRecovery,
+  reoptInitialDiscovery,
   reoptRecoveryLoop,
   resolveHeader,
   updateRecoveredModule,
@@ -1577,7 +1577,7 @@ findDebugDynDep opts depName = do
       )
 
 -- | Add the debug information from dynamic dependencies (if available either in
---   cached format or on disk for analysis).
+-- cached format or on disk for analysis).
 addDynDepDebugInfo ::
   ReoptOptions ->
   -- | Map to extend with debug info.
@@ -2345,7 +2345,7 @@ data RecoverX86Output = RecoverX86Output
   , summaryFailures :: Map (MemSegmentOff 64) (FunctionArgAnalysisFailure 64)
   }
 
--- | Analyze an elf binary to extract information.
+-- | Analyze an ELF binary to extract information.
 doRecoverX86 ::
   -- | Prefix to use if we need to generate new function endpoints later.
   BSC.ByteString ->
@@ -2616,13 +2616,15 @@ reoptX86Init loadOpts reoptOpts hdrInfo = do
       doInit loadOpts hdrInfo ainfo pltFn reoptOpts
   pure (os, initState)
 
-checkSymbolUnused ::
+-- | Checks that the prefix we intend to use for unnamed functions is not used
+-- by any of the pre-existing symbols.
+checkNoSymbolUsesReservedPrefix ::
   -- | Prefix to use if we need to generate new function endpoints later.
   BSC.ByteString ->
   -- | Symbol map constructor for binary.
   SymAddrMap (Macaw.ArchAddrWidth arch) ->
   ReoptM arch r ()
-checkSymbolUnused unnamedFunPrefix symAddrMap = do
+checkNoSymbolUsesReservedPrefix unnamedFunPrefix symAddrMap = do
   when (isUsedPrefix unnamedFunPrefix symAddrMap) $ do
     reoptFatalError $
       Events.ReoptInitError $
@@ -2649,13 +2651,15 @@ fnStmtHasCandidate (FnCall _fn args _mRet) = do
 fnStmtHasCandidate _ = return []
 
 -- | Repeatedly perform Macaw recovery and discover new potential function entry
--- points.
+-- points.  Incrementally re-runs discovery if new function entry points are
+-- identified.
 reoptRecoveryLoop ::
   SymAddrMap 64 ->
   ReoptOptions ->
   BSC.ByteString ->
   SyscallPersonality ->
   FunTypeMaps 64 ->
+  -- | Discovery state after the initial discovery run
   Macaw.DiscoveryState X86_64 ->
   ReoptM
     X86_64
@@ -2665,13 +2669,11 @@ reoptRecoveryLoop ::
     , RecoveredModule X86_64
     , ModuleConstraints X86_64
     )
-reoptRecoveryLoop symAddrMap rOpts funPrefix sysp debugTypeMap = go
+reoptRecoveryLoop symAddrMap rOpts funPrefix sysp debugTypeMap firstDiscState = do
+  checkNoSymbolUsesReservedPrefix funPrefix symAddrMap
+  go firstDiscState
  where
-  go previousDiscState = do
-    discState <-
-      reoptRunDiscovery (getAddrSymMap symAddrMap) $
-        Macaw.incCompleteDiscovery previousDiscState (roDiscoveryOptions rOpts)
-
+  go discState = do
     recoverX86Output <- doRecoverX86 funPrefix sysp symAddrMap debugTypeMap discState
     let recMod = recoveredModule recoverX86Output
 
@@ -2690,24 +2692,29 @@ reoptRecoveryLoop symAddrMap rOpts funPrefix sysp debugTypeMap = go
     -- NOTE: if we mark addresses that have already been tried (even if they
     -- have failed), Macaw will not add them to the unexplored frontier, so
     -- there is no risk here.
-    let newDiscState = Macaw.markAddrsAsFunction Macaw.UserRequest candidateAddressesAsSegOffs discState
-    let unexplored = newDiscState ^. Macaw.unexploredFunctions
+    let markedDiscState = Macaw.markAddrsAsFunction Macaw.UserRequest candidateAddressesAsSegOffs discState
+    let unexplored = markedDiscState ^. Macaw.unexploredFunctions
 
     if null unexplored
-      then traceM "NOLOOP" >> return (newDiscState, recoverX86Output, recMod, moduleConstraints)
-      else traceM "LOOP" >> go newDiscState
-
-reoptPrepareForRecovery ::
+      then return (markedDiscState, recoverX86Output, recMod, moduleConstraints)
+      else do
+        newDiscState <-
+          reoptRunDiscovery (getAddrSymMap symAddrMap) $
+            Macaw.incCompleteDiscovery markedDiscState (roDiscoveryOptions rOpts)
+        go newDiscState
+
+-- | Performs the first instance of Macaw discovery, before any recovery has
+-- happened.  This may yield less code than what a full recovery would, due to
+-- the discovery loop in `reoptRecoveryLoop`.
+reoptInitialDiscovery ::
   LoadOptions ->
   ReoptOptions ->
   AnnDeclarations ->
-  BSC.ByteString ->
   Elf.ElfHeaderInfo 64 ->
   ReoptM X86_64 r (X86OS, SymAddrMap 64, FunTypeMaps 64, Macaw.DiscoveryState X86_64)
-reoptPrepareForRecovery loadOpts reoptOpts hdrAnn unnamedFunPrefix hdrInfo = do
+reoptInitialDiscovery loadOpts reoptOpts hdrAnn hdrInfo = do
   (os, initState) <- reoptX86Init loadOpts reoptOpts hdrInfo
   let symAddrMap = initDiscSymAddrMap initState
-  checkSymbolUnused unnamedFunPrefix symAddrMap
   let ainfo = osArchitectureInfo os
   (debugTypeMap, discState) <- doDiscovery hdrAnn hdrInfo ainfo initState reoptOpts
   return (os, symAddrMap, debugTypeMap, discState)
@@ -2732,7 +2739,8 @@ recoverX86Elf ::
     , ModuleConstraints X86_64
     )
 recoverX86Elf loadOpts reoptOpts hdrAnn unnamedFunPrefix hdrInfo = do
-  (os, symAddrMap, debugTypeMap, discState) <- reoptPrepareForRecovery loadOpts reoptOpts hdrAnn unnamedFunPrefix hdrInfo
+  (os, symAddrMap, debugTypeMap, discState) <-
+    reoptInitialDiscovery loadOpts reoptOpts hdrAnn hdrInfo
   let sysp = osPersonality os
   (finalDiscState, recoverX86Output, recMod, moduleConstraints) <-
     reoptRecoveryLoop symAddrMap reoptOpts unnamedFunPrefix sysp debugTypeMap discState