Merge pull request #23 from nikitalita/fix-rebuilding-control-flow

Fix RebuildingControlFlow()

Champollion/main.cpp

@@ -30,10 +30,13 @@ struct Params
     bool dumpTree;
     bool recreateDirStructure;
     bool decompileDebugFuncs;
+    bool recursive;
 
     fs::path assemblyDir;
     fs::path papyrusDir;
 
+    fs::path parentDir{};
+
     std::vector<fs::path> inputs;
 };
 
@@ -62,6 +65,7 @@ OptionsResult getProgramOptions(int argc, char* argv[], Params& params)
             ("help,h", "Display the help message")
             ("asm,a", options::value<std::string>()->implicit_value(""), "If defined, output assembly file(s) to this directory")
             ("psc,p", options::value<std::string>(), "Name of the output dir for psc decompilation")
+            ("recursive,r", "Recursively scan directories for pex files")
             ("recreate-subdirs,s", "Recreates directory structure for script in root of output directory (Fallout 4 only, default false)")
             ("comment,c", "Output assembly in comments of the decompiled psc file")
             ("header,e", "Write header to decompiled psc file")
@@ -110,6 +114,7 @@ OptionsResult getProgramOptions(int argc, char* argv[], Params& params)
     params.parallel = (args.count("threaded") != 0);
     params.traceDecompilation = (args.count("trace") != 0);
     params.dumpTree = params.traceDecompilation && args.count("no-dump-tree") == 0;
+    params.recursive = (args.count("recursive") != 0);
     params.recreateDirStructure = (args.count("recreate-subdirs") != 0);
     params.decompileDebugFuncs = (args.count("debug-funcs") != 0);
     try
@@ -179,7 +184,13 @@ OptionsResult getProgramOptions(int argc, char* argv[], Params& params)
     return Good;
 }
 
-typedef std::vector<std::string> ProcessResults;
+struct _ProcessResults{
+    std::vector<std::string> output;
+    bool isStarfield = false;
+    bool failed = false;
+};
+
+typedef _ProcessResults ProcessResults;
 ProcessResults processFile(fs::path file, Params params)
 {
     ProcessResults result;
@@ -192,9 +203,11 @@ ProcessResults processFile(fs::path file, Params params)
     }
     catch(std::exception& ex)
     {
-       result.push_back(std::format("ERROR: {} : {}", file.string(), ex.what()));
+       result.output.push_back(std::format("ERROR: {} : {}", file.string(), ex.what()));
+       result.failed = true;
        return result;
     }
+    pex.getGameType() == Pex::Binary::StarfieldScript ? result.isStarfield = true : result.isStarfield = false;
     if (params.outputAssembly)
     {
         fs::path asmFile = params.assemblyDir / file.filename().replace_extension(".pas");
@@ -204,19 +217,22 @@ ProcessResults processFile(fs::path file, Params params)
             Decompiler::AsmCoder asmCoder(new Decompiler::StreamWriter(asmStream));
 
             asmCoder.code(pex);
-            result.push_back(std::format("{} dissassembled to {}", file.string(), asmFile.string()));
+            result.output.push_back(std::format("{} dissassembled to {}", file.string(), asmFile.string()));
         }
         catch(std::exception& ex)
         {
-            result.push_back(std::format("ERROR: {} : {}",file.string(),ex.what()));
+            result.output.push_back(std::format("ERROR: {} : {}", file.string(), ex.what()));
+            result.failed = true;
             fs::remove(asmFile);
         }
     }
-    fs::path dir_structure = "";
+    fs::path dir_structure;
     if (params.recreateDirStructure && (pex.getGameType() == Pex::Binary::Fallout4Script || pex.getGameType() == Pex::Binary::StarfieldScript) && pex.getObjects().size() > 0){
         std::string script_path = pex.getObjects()[0].getName().asString();
         std::replace(script_path.begin(), script_path.end(), ':', '/');
         dir_structure = fs::path(script_path).remove_filename();
+    } else if (!params.parentDir.empty()) {
+      dir_structure = fs::relative(file, params.parentDir).remove_filename();
     }
     fs::path basedir = !dir_structure.empty() ? (params.papyrusDir / dir_structure) : params.papyrusDir;
     if (!dir_structure.empty()){
@@ -240,22 +256,40 @@ ProcessResults processFile(fs::path file, Params params)
                 params.papyrusDir.string()); // using string instead of path here for C++14 compatability for staticlib targets
 
         pscCoder.code(pex);
-        result.push_back(std::format("{} decompiled to {}", file.string(), pscFile.string()));
+        result.output.push_back(std::format("{} decompiled to {}", file.string(), pscFile.string()));
     }
     catch(std::exception& ex)
     {
-        result.push_back(std::format("ERROR: {} : {}", file.string() , ex.what()));
+        result.output.push_back(std::format("ERROR: {} : {}", file.string() , ex.what()));
+        result.failed = true;
         fs::remove(pscFile);
     }
     return result;
 
 }
+size_t countFiles = 0;
+size_t failedFiles = 0;
+bool printStarfieldWarning = false;
+
+void processResult(ProcessResults result)
+{
+    if (result.failed){
+        ++failedFiles;
+    }
+    if (!printStarfieldWarning && result.isStarfield){
+      printStarfieldWarning = true;
+    }
+    for (auto line : result.output)
+    {
+        std::cout << line << '\n';
+    }
+}
+
 
 int main(int argc, char* argv[])
 {
 
     Params args;
-    size_t countFiles = 0;
     auto result = getProgramOptions(argc, argv, args);
     if (result == Good)
     {
@@ -264,30 +298,31 @@ int main(int argc, char* argv[])
         {
             for (auto path : args.inputs)
             {
-                if (fs::is_directory(path))
-                {
+                if (args.recursive && fs::is_directory(path)){
+                    args.parentDir = path;
+                    // recursively get all files in the directory
+                    for (auto& entry : fs::recursive_directory_iterator(path)){
+                        if (fs::is_regular_file(entry) && _stricmp(entry.path().extension().string().c_str(), ".pex") == 0){
+                            processResult(processFile(entry, args));
+                        }
+                    }
+                } else if (fs::is_directory(path)){
+                    args.parentDir = fs::path();
                     fs::directory_iterator end;
                     fs::directory_iterator entry(path);
                     while(entry != end)
                     {
                         if (_stricmp(entry->path().extension().string().c_str(), ".pex") == 0)
                         {
-                            for (auto line : processFile(entry->path(), args))
-                            {
-                                std::cout << line << '\n';
-                            }
-                            ++countFiles;
+                            processResult(processFile(path, args));
                         }
                         entry++;
                     }
                 }
                 else
                 {
-                    ++countFiles;
-                    for (auto line : processFile(path, args))
-                    {
-                        std::cout << line << '\n';
-                    }
+                  args.parentDir = fs::path();
+                  processResult(processFile(path, args));
                 }
             }
         }
@@ -296,8 +331,19 @@ int main(int argc, char* argv[])
             std::vector<std::future<ProcessResults>> results;
             for (auto& path : args.inputs)
             {
-                if (fs::is_directory(path))
+
+                if (args.recursive && fs::is_directory(path)){
+                  args.parentDir = path;
+                  // recursively get all files in the directory
+                  for (auto& entry : fs::recursive_directory_iterator(path)){
+                    if (fs::is_regular_file(entry) && _stricmp(entry.path().extension().string().c_str(), ".pex") == 0){
+                        results.push_back(std::move(std::async(std::launch::async, processFile, fs::path(entry.path()), args)));
+                    }
+                  }
+                }
+                else if (fs::is_directory(path))
                 {
+                    args.parentDir = fs::path();
                     fs::directory_iterator end;
                     fs::directory_iterator entry(path);
                     while(entry != end)
@@ -312,16 +358,25 @@ int main(int argc, char* argv[])
                 }
                 else
                 {
+                    args.parentDir = fs::path();
                     results.push_back(std::move(std::async(std::launch::async, processFile, path, args)));
                 }
             }
 
             for (auto& result : results)
             {
-                for(auto& line : result.get())
+                auto processResult = result.get();
+                if (!printStarfieldWarning && processResult.isStarfield){
+                  printStarfieldWarning = true;
+                }
+
+                for(auto& line : processResult.output)
                 {
                     std::cout << line << '\n';
                 }
+                if (processResult.failed){
+                    ++failedFiles;
+                }
             }
             countFiles = results.size();
 
@@ -330,6 +385,23 @@ int main(int argc, char* argv[])
         auto diff = end - start;
 
         std::cout << countFiles << " files processed in " << std::chrono::duration <double> (diff).count() << " s" << std::endl;
+        if (failedFiles > 0){
+            std::cout << failedFiles << " files failed to decompile." << std::endl;
+        }
+
+        if (printStarfieldWarning){
+          // TODO: Remove this warning when the CK comes out
+          std::cout << "********************* STARFIELD PRELIMINARY SYNTAX WARNING *********************" << std::endl;
+          std::cout << "The syntax for new features in Starfield (Guard, TryGuard, GetMatchingStructs) is not yet known." << std::endl;
+          std::cout << "Decompiled Starfield scripts use guessed-at syntax for these features." << std::endl;
+          std::cout << "This syntax should be considered as experimental, unstable, and subject to change." << std::endl << std::endl;
+          std::cout << "The proper syntax will only be known when the Creation Kit comes out in early 2024." << std::endl;
+          std::cout << "If you are using decompiled scripts as the basis for mods, please be aware of this," << std::endl;
+          std::cout << "and be prepared to update your scripts when the final syntax is known." << std::endl << std::endl;
+          std::cout << "The lines in the decompiled scripts which contain this guessed-at syntax are" << std::endl;
+          std::cout << "marked with a comment beginning with '" << Decompiler::WARNING_COMMENT_PREFIX << "'." << std::endl;
+          std::cout << "********************* STARFIELD PRELIMINARY SYNTAX WARNING *********************" << std::endl;
+        }
         return 0;
     }
     if (result == HelpOrVersion){

Decompiler/PscCodeGenerator.cpp

@@ -32,7 +32,7 @@ Decompiler::PscCodeGenerator::PscCodeGenerator(Decompiler::PscDecompiler* decomp
 void Decompiler::PscCodeGenerator::newLine()
 {
     if (!m_ExperimentalSyntaxWarning.empty()) {
-        m_Result << " ;*** WARNING: Experimental syntax, may be incorrect: ";
+        m_Result << " " << Decompiler::WARNING_COMMENT_PREFIX << " WARNING: Experimental syntax, may be incorrect: ";
         for (auto warn: m_ExperimentalSyntaxWarning){
             m_Result << warn << " ";
         }

Decompiler/PscCoder.cpp

@@ -216,7 +216,7 @@ void Decompiler::PscCoder::writeObject(const Pex::Object &object, const Pex::Bin
     if (object.getGuards().size()) {
         write("");
         write(";-- Guards ------------------------------------------");
-        write(";*** WARNING: Guard declaration syntax is EXPERIMENTAL, subject to change");
+        write(std::string(Decompiler::WARNING_COMMENT_PREFIX) + " WARNING: Guard declaration syntax is EXPERIMENTAL, subject to change");
         writeGuards(object, pex);
     }
 

Decompiler/PscCoder.hpp

@@ -3,6 +3,7 @@
 
 
 namespace Decompiler {
+static const char* WARNING_COMMENT_PREFIX = ";***";
 /**
  * @brief Write a PEX file as a PSC file.
  */

Decompiler/PscDecompiler.cpp

@@ -479,7 +479,9 @@ void Decompiler::PscDecompiler::createNodesForBlocks(size_t block)
                         node = std::make_shared<Node::Cast>(ip, args[0].getId(), fromValue(ip, args[1]), typeOfVar(args[0].getId()));
                     } else // two variables of the same type, equivalent to an assign
                     {
-                        node = std::make_shared<Node::Copy>(ip, args[0].getId(), fromValue(ip, args[1]));
+                        // check if this is a useless cast
+                        if (args[0].getId() != args[1].getId())
+                          node = std::make_shared<Node::Copy>(ip, args[0].getId(), fromValue(ip, args[1]));
                     }
                     break;
                 }
@@ -826,7 +828,8 @@ void Decompiler::PscDecompiler::rebuildExpression(Node::BasePtr scope)
             }
             else
             {
-                throw std::runtime_error("Decompilation failed");
+              auto funcname = m_Function.getName().isValid() ? m_Function.getName().asString() : "unknown function";
+              throw std::runtime_error("Failed to rebuild expression in " + funcname + " at instruction " + std::to_string(expressionUse->getBegin()));
             }
         }
         else
@@ -1015,7 +1018,8 @@ Node::BasePtr Decompiler::PscDecompiler::rebuildControlFlow(size_t startBlock, s
 {
     if (endBlock < startBlock)
     {
-        throw std::runtime_error("Decompilation failed");
+      auto funcname = m_Function.getName().isValid() ? m_Function.getName().asString() : "unknown function";
+      throw std::runtime_error("Failed to rebuild control flow for " + funcname + ".");
     }
     auto begin = m_CodeBlocs.find(startBlock);
     auto end = m_CodeBlocs.find(endBlock);
@@ -1036,9 +1040,20 @@ Node::BasePtr Decompiler::PscDecompiler::rebuildControlFlow(size_t startBlock, s
             if (beforeExit == PscCodeBlock::END)
             {
                 // Decompilation failed
-                throw std::runtime_error("Decompilation failed");
+                auto funcname = m_Function.getName().isValid() ? m_Function.getName().asString() : "unknown function";
+                throw std::runtime_error("Failed to rebuild control flow for " + funcname + ".");
             }
-            auto& lastBlock = m_CodeBlocs[beforeExit];
+            Node::BasePtr condition = std::make_shared<Node::Constant>(-1, Pex::Value(source->getCondition(), true));
+
+            if (m_CodeBlocs[beforeExit] == source){
+                assert(source->onFalse() < source->onTrue());
+                auto scope = source->getScope();
+                condition = std::make_shared<Node::UnaryOperator>(-1, 10, source->getCondition(), "!", condition);
+                source->setCondition(source->getCondition(), source->onFalse(), source->onTrue());
+                exit = source->onFalse();
+                beforeExit = findBlockForInstruction(exit-1);
+              }
+              auto& lastBlock = m_CodeBlocs[beforeExit];
 
             // The last block is an unconditional jump to the current block
             // This is a while.
@@ -1048,15 +1063,12 @@ Node::BasePtr Decompiler::PscDecompiler::rebuildControlFlow(size_t startBlock, s
                 auto whileStartBlock = source->onTrue();
                 auto whileEndBlock = source->onFalse();
 
-
-                Node::BasePtr whileCondition = std::make_shared<Node::Constant>(-1, Pex::Value(source->getCondition(), true));
-
                 result->mergeChildren(source->getScope()->shared_from_this());
 
                 // Rebuild the statements in the while loop.
                 auto whileBody = rebuildControlFlow(whileStartBlock, whileEndBlock);
 
-                *result << std::make_shared<Node::While>(-1, whileCondition, whileBody);
+                *result << std::make_shared<Node::While>(-1, condition, whileBody);
 
                 advance = 0;
                 it = m_CodeBlocs.find(whileEndBlock);
@@ -1071,14 +1083,12 @@ Node::BasePtr Decompiler::PscDecompiler::rebuildControlFlow(size_t startBlock, s
                     auto ifStartBlock = source->onTrue();
                     auto ifEndBlock = source->onFalse();
 
-                    Node::BasePtr ifCondition = std::make_shared<Node::Constant>(-1, Pex::Value(source->getCondition(), true));
-
                     result->mergeChildren(source->getScope()->shared_from_this());
 
                     // Rebuild the statements of the if body
                     auto ifBody = rebuildControlFlow(ifStartBlock, ifEndBlock);
 
-                    *result << std::make_shared<Node::IfElse>(-1, ifCondition, ifBody, nullptr);
+                    *result << std::make_shared<Node::IfElse>(-1, condition, ifBody, nullptr);
 
                     advance = 0;
                     it = m_CodeBlocs.find(ifEndBlock);
@@ -1089,7 +1099,6 @@ Node::BasePtr Decompiler::PscDecompiler::rebuildControlFlow(size_t startBlock, s
                     auto elseStartBlock = source->onFalse();
                     auto endElseBlock = lastBlock->getNext();
 
-                    Node::BasePtr ifCondition = std::make_shared<Node::Constant>(-1, Pex::Value(source->getCondition(), true));
 
                     result->mergeChildren(source->getScope()->shared_from_this());
 
@@ -1098,7 +1107,7 @@ Node::BasePtr Decompiler::PscDecompiler::rebuildControlFlow(size_t startBlock, s
                     // Rebuilds the statements in the else body.
                     auto elseBody = rebuildControlFlow(elseStartBlock, endElseBlock);
 
-                    *result << std::make_shared<Node::IfElse>(-1, ifCondition, ifBody, elseBody);
+                    *result << std::make_shared<Node::IfElse>(-1, condition, ifBody, elseBody);
 
                     advance = 0;
                     it = m_CodeBlocs.find(endElseBlock);
@@ -1376,6 +1385,17 @@ void Decompiler::PscDecompiler::cleanUpTree(Node::BasePtr program)
 
     program->computeInstructionBounds();
 
+    // check for orphaned nodes on m_CodeBlocs
+    for (auto& bloc_kv : m_CodeBlocs)
+    {
+        auto& bloc = bloc_kv.second;
+        auto scope = bloc->getScope();
+        if (scope->size() > 0) {
+          auto funcname = m_Function.getName().isValid() ? m_Function.getName().asString() : "unknown function";
+          throw std::runtime_error("Orphaned nodes in " + funcname + " from instruction " + std::to_string(scope->front()->getBegin()) + " to " + std::to_string(scope->back()->getEnd()) + ".");
+        }
+    }
+
 }
 
 

README.md

@@ -16,6 +16,7 @@ Champollion is a CLI-only program.
 | -a [*assembly directory*] | --asm [*assembly directory*] | Champollion will write an assembly version of the PEX file in the given directory, if one. The assembly file is an human readable version of the content of the PEX file |
 | -c                        | --comment                    | The decompiled file will be annotated with the assembly instruction corresponding to the decompiled code lines. |
 | -t                        | --threaded                   | Champollion will parallelize the decompilation. It is useful when decompiling a directory containing many PEX files. |
+| -r                        | --recursive                  | Recursively scan specified directory(s) for pex files to decompile|
 | -s                        | --recreate-subdirs           | Recreates directory structure for script in root of output directory (Fallout 4 only, default false) |
 | -e                        | --header                     | Write header to decompiled psc file                          |
 | -g                        | --trace                      | Trace the decompilation and output results to rebuild log    |