codegen: refactor and cleanup

CMakeLists.txt

@@ -1,5 +1,4 @@
 # Copyright (c) 2024, Trail of Bits, Inc.
-#
 # This source code is licensed in accordance with the terms specified in the
 # LICENSE file found in the root directory of this source tree.
 

include/patchestry/AST/Codegen.hpp

@@ -7,8 +7,11 @@
 
 #pragma once
 
+#include <unordered_map>
+
 #include <clang/AST/ASTContext.h>
 #include <clang/Frontend/CompilerInstance.h>
+
 #include <vast/Frontend/FrontendAction.hpp>
 #include <vast/Frontend/Options.hpp>
 
@@ -17,6 +20,9 @@ namespace llvm {
 }
 
 namespace patchestry::ast {
+
+    using LocationMap = std::unordered_map< void *, std::string >;
+
     class CodeGenerator
     {
       public:
@@ -29,7 +35,9 @@ namespace patchestry::ast {
 
         virtual ~CodeGenerator() {}
 
-        void generate_source_ir(clang::ASTContext &ctx, llvm::raw_fd_ostream &os);
+        void generate_source_ir(
+            clang::ASTContext &ctx, const LocationMap &locations, llvm::raw_fd_ostream &os
+        );
 
       private:
         vast::cc::action_options opts;

include/patchestry/AST/FunctionBuilder.hpp

@@ -7,17 +7,18 @@
 
 #pragma once
 
-#include "patchestry/AST/TypeBuilder.hpp"
+#include <functional>
+#include <memory>
+#include <unordered_map>
+
+#include <clang/AST/ASTContext.h>
 #include <clang/AST/Decl.h>
 #include <clang/AST/DeclBase.h>
 #include <clang/Frontend/CompilerInstance.h>
 #include <clang/Sema/Sema.h>
-#include <functional>
 
-#include <clang/AST/ASTContext.h>
-#include <memory>
+#include <patchestry/AST/TypeBuilder.hpp>
 #include <patchestry/Ghidra/JsonDeserialize.hpp>
-#include <unordered_map>
 
 namespace patchestry::ast {
     class OpBuilder;
@@ -88,6 +89,13 @@ namespace patchestry::ast {
 
         clang::FunctionDecl *create_definition(clang::ASTContext &ctx);
 
+        template< typename T >
+        void set_location_key(T *pointer, const std::string &key) {
+            if (!location_map.get().contains(pointer)) {
+                location_map.get().emplace(pointer, key);
+            }
+        }
+
       private:
         void create_labels(clang::ASTContext &ctx, clang::FunctionDecl *func_decl);
 

include/patchestry/AST/OperationBuilder.hpp

@@ -154,7 +154,7 @@ namespace patchestry::ast {
       private:
         clang::Stmt *create_assign_operation(
             clang::ASTContext &ctx, clang::Expr *input_expr, clang::Expr *output_expr,
-            clang::SourceLocation location
+            const std::string &location_key
         );
 
         /**
@@ -172,8 +172,10 @@ namespace patchestry::ast {
             clang::ASTContext &ctx, clang::Expr *expr, clang::QualType to_type
         );
 
-        clang::Stmt *
-        create_varnode(clang::ASTContext &ctx, const Function &function, const Varnode &vnode);
+        clang::Stmt *create_varnode(
+            clang::ASTContext &ctx, const Function &function, const Varnode &vnode,
+            const std::string &op_key = ""
+        );
 
         clang::Stmt *create_parameter(clang::ASTContext &ctx, const Varnode &vnode);
 

include/patchestry/AST/TypeBuilder.hpp

@@ -7,8 +7,10 @@
 
 #pragma once
 
-#include <clang/AST/ASTContext.h>
 #include <functional>
+
+#include <clang/AST/ASTContext.h>
+
 #include <patchestry/Ghidra/JsonDeserialize.hpp>
 
 namespace patchestry::ast {

include/patchestry/Dialect/Pcode/Deserialize.hpp

@@ -15,12 +15,6 @@
 
 namespace patchestry::pc {
 
-    struct program;
-    struct function;
-    struct basic_block;
-    struct instruction;
-    struct pcode;
-
     using json_arr = llvm::json::Array;
     using json_obj = llvm::json::Object;
     using json_val = llvm::json::Value;
@@ -36,14 +30,10 @@ namespace patchestry::pc {
             bld.setInsertionPointToStart(&*reg.begin());
         }
 
-        void process(const program &prog);
-        void process_function(const function &func);
-        void process_block(const basic_block &block);
-        void process_instruction(const instruction &inst);
-        void process_pcode(const pcode &code);
-
-        mlir_operation create_int_const(uint32_t offset, uint32_t size);
-        mlir_operation create_varnode(std::string type, uint32_t offset, uint32_t size);
+        void process(const json_obj &json);
+        void process_function(const json_obj &json);
+        void process_block(const json_obj &json);
+        void process_instruction(const json_obj &json);
     };
 
     mlir::OwningOpRef< mlir::ModuleOp > deserialize(const json_obj &json, mcontext_t *mctx);