Proper syntax generation for Lock/TryLock

Orvid · Sep 11, 2023 · 2881df1 · 2881df1
1 parent e4a4b5f
commit 2881df1
Show file tree

Hide file tree

Showing 14 changed files with 377 additions and 23 deletions.
diff --git a/Decompiler/DumpTree.hpp b/Decompiler/DumpTree.hpp
@@ -268,6 +268,41 @@ class DumpTree : public Node::VisitorBase
         node->getObject()->visit(this);
         leave();
     }
+
+    virtual void visit(Node::Lock* node)
+    {
+        enter("Lock", node);
+        for (auto param : *node->getParameters())
+        {
+            m_Callback(indent() << "Param:");
+            param->visit(this);
+        }
+        m_Callback(indent() << "Body:");
+        node->getBody()->visit(this);
+        leave();
+    }
+    virtual void visit(Node::TryLock* node)
+    {
+        enter("Lock", node);
+        for (auto param : *node->getParameters())
+        {
+            m_Callback(indent() << "Param:");
+            param->visit(this);
+        }
+        m_Callback(indent() << "Body:");
+        node->getBody()->visit(this);
+        leave();
+    }
+    virtual void visit(Node::Unlock* node)
+    {
+        enter("Unlock", node);
+        for (auto param : *node->getParameters())
+        {
+            m_Callback(indent() << "Param:");
+            param->visit(this);
+        }
+        leave();
+    }
 };
 
 }
diff --git a/Decompiler/Node/Lock.hpp b/Decompiler/Node/Lock.hpp
@@ -0,0 +1,40 @@
+#pragma once
+
+#include <cassert>
+#include <cstdint>
+#include <memory>
+
+#include "Base.hpp"
+#include "FieldNodeMixin.hpp"
+#include "Scope.hpp"
+#include "Visitor.hpp"
+
+namespace Node {
+
+    class Lock final :
+            public Base,
+            public FieldParametersNodeMixin<0>,
+            public FieldBodyNodeMixin<1>
+    {
+    public:
+        Lock(size_t ip, BasePtr body) :
+                Base(2, ip, 10),
+                FieldParametersNodeMixin(this, std::make_shared<Params>()),
+                FieldBodyNodeMixin(this, body)
+        {
+        }
+        virtual ~Lock() = default;
+
+        void visit(Visitor* visitor) override
+        {
+            assert(visitor);
+            visitor->visit(this);
+        }
+
+        void computeInstructionBounds() override
+        {
+            Base::computeInstructionBounds();
+        }
+    };
+
+}
diff --git a/Decompiler/Node/NodeComparer.cpp b/Decompiler/Node/NodeComparer.cpp
@@ -182,7 +182,6 @@ void Node::NodeComparer::visit(Node::While *node)
     };
     m_Reference->visit(& m_Visitor);
 }
-
 void Node::NodeComparer::visit(Node::IfElse *node)
 {
     m_Visitor << (DynamicVisitor::LambdaIfElse)[&](IfElse* ref, DynamicVisitor*) {
@@ -199,6 +198,30 @@ void Node::NodeComparer::visit(Node::Declare *node)
     m_Reference->visit(& m_Visitor);
 }
 
+void Node::NodeComparer::visit(Node::Lock *node)
+{
+    m_Visitor << (DynamicVisitor::LambdaLock)[&](Lock* ref, DynamicVisitor*) {
+        m_Result = isSameChildren(ref, node);
+    };
+    m_Reference->visit(& m_Visitor);
+}
+
+void Node::NodeComparer::visit(Node::TryLock *node)
+{
+    m_Visitor << (DynamicVisitor::LambdaLock)[&](Lock* ref, DynamicVisitor*) {
+        m_Result = isSameChildren(ref, node);
+    };
+    m_Reference->visit(& m_Visitor);
+}
+
+void Node::NodeComparer::visit(Node::Unlock *node)
+{
+    m_Visitor << (DynamicVisitor::LambdaLock)[&](Lock* ref, DynamicVisitor*) {
+        m_Result = isSameChildren(ref, node);
+    };
+    m_Reference->visit(& m_Visitor);
+}
+
 bool Node::isSameTree(Node::BasePtr left, Node::BasePtr right)
 {
     NodeComparer comparer(left);

diff --git a/Decompiler/Node/NodeComparer.hpp b/Decompiler/Node/NodeComparer.hpp
@@ -39,11 +39,15 @@ class NodeComparer final : public Visitor
     void visit(Node::While* node) override;
     void visit(Node::IfElse* node) override;
     void visit(Node::Declare* node) override;
+    void visit(Node::Lock *node) override;
+    void visit(Node::TryLock *node) override;
 
 private:
     BasePtr m_Reference;
     DynamicVisitor m_Visitor;
     bool m_Result;
+
+    void visit(Unlock *node);
 };
 
 bool isSameTree(BasePtr left, BasePtr right);

diff --git a/Decompiler/Node/Nodes.hpp b/Decompiler/Node/Nodes.hpp
@@ -31,4 +31,9 @@
 
 #include "Declare.hpp"
 
+#include "Lock.hpp"
+#include "TryLock.hpp"
+#include "Unlock.hpp"
+
+
 #endif // NODES_HPP
diff --git a/Decompiler/Node/TryLock.hpp b/Decompiler/Node/TryLock.hpp
@@ -0,0 +1,40 @@
+#pragma once
+
+#include <cassert>
+#include <cstdint>
+#include <memory>
+
+#include "Base.hpp"
+#include "FieldNodeMixin.hpp"
+#include "Scope.hpp"
+#include "Visitor.hpp"
+
+namespace Node {
+
+    class TryLock final :
+            public Base,
+            public FieldParametersNodeMixin<0>,
+            public FieldBodyNodeMixin<1>
+    {
+    public:
+        TryLock(size_t ip, const Pex::StringTable::Index& result, BasePtr body) :
+                Base(2, ip, 10, result),
+                FieldParametersNodeMixin(this, std::make_shared<Params>()),
+                FieldBodyNodeMixin(this, body)
+        {
+        }
+        virtual ~TryLock() = default;
+
+        void visit(Visitor* visitor) override
+        {
+            assert(visitor);
+            visitor->visit(this);
+        }
+
+        void computeInstructionBounds() override
+        {
+            Base::computeInstructionBounds();
+        }
+    };
+
+}
diff --git a/Decompiler/Node/Unlock.hpp b/Decompiler/Node/Unlock.hpp
@@ -0,0 +1,38 @@
+#pragma once
+
+#include <cassert>
+#include <cstdint>
+#include <memory>
+
+#include "Base.hpp"
+#include "FieldNodeMixin.hpp"
+#include "Scope.hpp"
+#include "Visitor.hpp"
+
+namespace Node {
+
+    class Unlock final :
+            public Base,
+            public FieldParametersNodeMixin<0>
+    {
+    public:
+        Unlock(size_t ip) :
+                Base(1, ip, 10),
+                FieldParametersNodeMixin(this, std::make_shared<Params>())
+        {
+        }
+        virtual ~Unlock() = default;
+
+        void visit(Visitor* visitor) override
+        {
+            assert(visitor);
+            visitor->visit(this);
+        }
+
+        void computeInstructionBounds() override
+        {
+            Base::computeInstructionBounds();
+        }
+    };
+
+}
diff --git a/Decompiler/Node/Visitor.hpp b/Decompiler/Node/Visitor.hpp
@@ -27,7 +27,10 @@ namespace Node {
     DO_NODE(While) \
     DO_NODE(IfElse) \
     DO_NODE(Declare) \
-    DO_NODE(StructCreate) \
+    DO_NODE(StructCreate)     \
+    DO_NODE(Lock)             \
+    DO_NODE(TryLock)          \
+    DO_NODE(Unlock)
 
 #define DO_NODE(NODE) class NODE;
 FOR_EACH_NODE_CLASS()

diff --git a/Decompiler/PscCodeBlock.cpp b/Decompiler/PscCodeBlock.cpp
@@ -148,3 +148,11 @@ Node::Scope *Decompiler::PscCodeBlock::getScope() const
 {
     return static_cast<Node::Scope*>(m_Scope.get());
 }
+
+void Decompiler::PscCodeBlock::addLockGuard(Pex::StringTable::Index index) {
+    m_Guards.push_back(index);
+}
+
+bool Decompiler::PscCodeBlock::isLock() const {
+    return !m_Guards.empty();
+}
diff --git a/Decompiler/PscCodeBlock.hpp b/Decompiler/PscCodeBlock.hpp
@@ -27,6 +27,8 @@ class PscCodeBlock
 public:
     static const size_t END = 0xFFFFFFFF;
 
+    void addLockGuard(Pex::StringTable::Index index);
+
 public:
     PscCodeBlock(size_t getBegin, size_t getEnd);
     ~PscCodeBlock();
@@ -40,6 +42,7 @@ class PscCodeBlock
     const Pex::StringTable::Index& getCondition() const;
 
     bool isConditional() const;
+    bool isLock() const;
 
     void setEnd(size_t getEnd);
     void setNext(size_t getNext);
@@ -56,6 +59,7 @@ class PscCodeBlock
     size_t m_Next;
     size_t m_OnFalse;
     Pex::StringTable::Index m_Condition;
+    std::vector<Pex::StringTable::Index> m_Guards;
 
     Node::BasePtr m_Scope;
 

diff --git a/Decompiler/PscCodeGenerator.cpp b/Decompiler/PscCodeGenerator.cpp
@@ -297,6 +297,7 @@ void Decompiler::PscCodeGenerator::visit(Node::While* node)
 
 void Decompiler::PscCodeGenerator::visit(Node::IfElse* node)
 {
+    auto cond = node->getCondition();
     m_Result << "If ";
     node->getCondition()->visit(this);
     m_Level++;
@@ -338,3 +339,30 @@ void Decompiler::PscCodeGenerator::visit(Node::Declare *node)
     node->getObject()->visit(this);
 }
 
+void Decompiler::PscCodeGenerator::visit(Node::Lock *node) {
+    m_Result << "Lock ";
+    node->getParameters()->visit(this);
+    m_Level++;
+    newLine();
+    node->getBody()->visit(this);
+    m_Level--;
+    newLine();
+    m_Result << "EndLock";
+}
+
+void Decompiler::PscCodeGenerator::visit(Node::TryLock *node) {
+    m_Result << "TryLock ";
+    node->getParameters()->visit(this);
+    m_Level++;
+    newLine();
+    node->getBody()->visit(this);
+    m_Level--;
+    newLine();
+    m_Result << "EndLock";
+
+}
+
+void Decompiler::PscCodeGenerator::visit(Node::Unlock *node) {
+    // NONE
+}
+
diff --git a/Decompiler/PscCodeGenerator.hpp b/Decompiler/PscCodeGenerator.hpp
@@ -42,6 +42,10 @@ class PscCodeGenerator : public Node::Visitor
     virtual void visit(Node::While* node);
     virtual void visit(Node::IfElse* node);
     virtual void visit(Node::Declare* node);
+    virtual void visit(Node::Lock* node);
+    virtual void visit(Node::TryLock* node);
+    virtual void visit(Node::Unlock* node);
+
 
 protected:    
     void newLine();

diff --git a/Decompiler/PscCoder.cpp b/Decompiler/PscCoder.cpp
@@ -589,6 +589,9 @@ void Decompiler::PscCoder::writeFunction(int i, const Pex::Function &function, c
         write(stream.str());
         writeDocString(i, function);
     } else {
+        if (functionName == "TaskmasterPossibleRestore"){
+            int i = 0; // TODO: REMOVE THIS HACK
+        }
         auto decomp = PscDecompiler(function, object, m_CommentAsm, m_TraceDecompilation, m_DumpTree,
                                     m_OutputDir);
         if (decomp.isDebugFunction()) {