git.seregonwar.com

core/decompiler/analysis/MemberAccessAnalysis.h

ShadPKG / core / decompiler / analysis / MemberAccessAnalysis.h

1	#pragma once
2
3	#include "../ir/AST.h"
4	#include "TypeSystem.h"
5	#include <memory>
6	#include <sstream>
7	#include <iomanip>
8
9	namespace ShadPKG::Decompiler::Analysis {
10
11	class MemberAccessAnalysis : public AST::ASTVisitor {
12	public:
13	explicit MemberAccessAnalysis(std::shared_ptr<TypeManager> typeManager)
14	: typeManager_(typeManager) {}
15
16	void analyze(std::shared_ptr<AST::FunctionAST> func) {
17	if (!func \|\| !func->body)
18	return;
19	func->body->accept(this);
20	}
21
22	// Visitor Implementation
23	void visit(AST::ConstantExpr *node) override {}
24	void visit(AST::VariableExpr *node) override {}
25
26	void visit(AST::BinaryExpr *node) override {
27	node->left = transform(node->left);
28	node->right = transform(node->right);
29	}
30
31	void visit(AST::UnaryExpr *node) override {
32	node->operand = transform(node->operand);
33	}
34
35	void visit(AST::CallExpr *node) override {
36	for (auto &arg : node->arguments)
37	arg = transform(arg);
38	}
39
40	void visit(AST::MemoryExpr *node) override {}
41
42	void visit(AST::MemberAccessExpr *node) override { node->base->accept(this); }
43
44	void visit(AST::CastExpr *node) override { node->expr = transform(node->expr); }
45
46	void visit(AST::CompoundStatement *node) override {
47	for (auto &stmt : node->statements)
48	stmt->accept(this);
49	}
50
51	void visit(AST::ExpressionStatement *node) override {
52	// We need to be able to replace node->expression if it changes.
53	// Let's call a helper "transform" that returns the new expression.
54	node->expression = transform(node->expression);
55	}
56
57	void visit(AST::IfStatement *node) override {
58	node->condition = transform(node->condition);
59	if (node->thenBranch)
60	node->thenBranch->accept(this);
61	if (node->elseBranch)
62	node->elseBranch->accept(this);
63	}
64
65	void visit(AST::WhileStatement *node) override {
66	node->condition = transform(node->condition);
67	if (node->body)
68	node->body->accept(this);
69	}
70
71	void visit(AST::DoWhileStatement *node) override {
72	node->condition = transform(node->condition);
73	if (node->body)
74	node->body->accept(this);
75	}
76
77	void visit(AST::ForStatement *node) override {} // Unimplemented
78
79	void visit(AST::ReturnStatement *node) override {
80	if (node->value)
81	node->value = transform(node->value);
82	}
83
84	void visit(AST::BreakStatement *node) override {}
85	void visit(AST::ContinueStatement *node) override {}
86	void visit(AST::GotoStatement *node) override {}
87	void visit(AST::LabelStatement *node) override {}
88	void visit(AST::CaseStmt *node) override {
89	// Switch bodies (CaseStmt) contain a CompoundStatement body
90	if (node->body)
91	node->body->accept(this);
92	}
93	void visit(AST::SwitchStmt *node) override {
94	node->condition = transform(node->condition);
95	for (auto &c : node->cases)
96	c->accept(this);
97	}
98
99	private:
100	std::shared_ptr<TypeManager> typeManager_;
101
102	std::shared_ptr<AST::Expression>
103	transform(std::shared_ptr<AST::Expression> expr) {
104	if (!expr)
105	return nullptr;
106
107	// Recurse children first? Or handle this node?
108	// Let's handle bottom-up in general, but here top-down might be enough.
109
110	// 1. Check if expr is MemoryExpr (dereference)
111	if (auto mem = std::dynamic_pointer_cast<AST::MemoryExpr>(expr)) {
112	// Check content of memory access
113	// Usually: (base + offset) or (base)
114
115	// We need to look inside mem->address (Expression)
116	// It is likely a BinaryExpr (ADD) or VariableExpr
117
118	/*
119	Case 1: *(ptr) -> ptr->member (at offset 0)
120	Case 2: *(ptr + 4) -> ptr->member (at offset 4)
121	*/
122
123	// Case 1: *base (offset 0) -> mem->offset is null
124	// Case 2: base[offset] -> mem->offset is set
125
126	std::shared_ptr<AST::Expression> base = mem->base;
127	int offset = 0;
128
129	if (mem->offset) {
130	if (auto cnst =
131	std::dynamic_pointer_cast<AST::ConstantExpr>(mem->offset)) {
132	offset = (int)cnst->intValue;
133	} else {
134	// Dynamic offset, cannot resolve to static member access easily
135	return mem; // Just enable recursion on base? No, return as is.
136	}
137	}
138
139	if (base) {
140	// Infer Struct
141	if (auto var = std::dynamic_pointer_cast<AST::VariableExpr>(base)) {
142	// DO NOT create structs for RIP-relative access
143	if (var->name == "rip") return mem;
144
145	// Create struct for this variable
146	std::string structName = "Struct_" + var->name;
147
148	auto type = typeManager_->getType(structName);
149	std::shared_ptr<StructType> structType;
150	if (!type) {
151	structType = typeManager_->createStruct(structName);
152	} else {
153	structType = std::dynamic_pointer_cast<StructType>(type);
154	}
155
156	if (structType) {
157	// Check if member exists
158	const StructMember* member = structType->getMemberAt(offset);
159	if (!member) {
160	std::stringstream ss;
161	ss << "field_0x" << std::hex << offset;
162	structType->addMember(ss.str(), typeManager_->getType("int"), offset);
163	member = structType->getMemberAt(offset);
164	}
165
166	// Transform this MemoryExpr into MemberAccessExpr
167	// base->field
168	return std::make_shared<AST::MemberAccessExpr>(
169	base,
170	member->name,
171	offset);
172	}
173	}
174	}
175	}
176
177	// Recurse for children if not replaced
178	// This 'transform' helper is manual recursion.
179	// Existing Visitor pattern doesn't support easy rewrite.
180	// I'll implement a simple Recursive Rewriter logic here for specific nodes.
181
182	if (auto bin = std::dynamic_pointer_cast<AST::BinaryExpr>(expr)) {
183	bin->left = transform(bin->left);
184	bin->right = transform(bin->right);
185	}
186	// ... (other recursive calls)
187
188	return expr;
189	}
190	};
191
192	} // namespace ShadPKG::Decompiler::Analysis
193

Seregon/ShadPKG