Skip to main content

thalir_core/
values.rs

1use crate::types::Type;
2use num_bigint::{BigInt, BigUint};
3use num_traits::ToPrimitive;
4use serde::{Deserialize, Serialize};
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
7pub enum ValueId {
8    Var(VarId),
9    Temp(TempId),
10    Param(ParamId),
11    BlockParam(BlockParamId),
12    Storage(StorageRefId),
13    Memory(MemoryRefId),
14    Global(GlobalId),
15}
16
17#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
18pub enum Value {
19    Register(ValueId),
20    Variable(VarId),
21    Temp(TempId),
22    Param(ParamId),
23    BlockParam(BlockParamId),
24    Constant(Constant),
25    StorageRef(StorageRefId),
26    MemoryRef(MemoryRefId),
27    Global(GlobalId),
28    Undefined,
29}
30
31impl Value {
32    pub fn as_register(&self) -> Option<ValueId> {
33        match self {
34            Value::Register(id) => Some(*id),
35            Value::Variable(v) => Some(ValueId::Var(*v)),
36            Value::Temp(t) => Some(ValueId::Temp(*t)),
37            Value::Param(p) => Some(ValueId::Param(*p)),
38            Value::BlockParam(bp) => Some(ValueId::BlockParam(*bp)),
39            Value::StorageRef(s) => Some(ValueId::Storage(*s)),
40            Value::MemoryRef(m) => Some(ValueId::Memory(*m)),
41            Value::Global(g) => Some(ValueId::Global(*g)),
42            _ => None,
43        }
44    }
45
46    pub fn is_constant(&self) -> bool {
47        matches!(self, Value::Constant(_))
48    }
49
50    pub fn is_reference(&self) -> bool {
51        matches!(self, Value::StorageRef(_) | Value::MemoryRef(_))
52    }
53
54    pub fn as_constant(&self) -> Option<&Constant> {
55        match self {
56            Value::Constant(c) => Some(c),
57            _ => None,
58        }
59    }
60}
61
62#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
63pub struct VarId(pub u32);
64
65impl std::fmt::Display for VarId {
66    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67        write!(f, "v{}", self.0)
68    }
69}
70
71#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
72pub struct TempId(pub u32);
73
74impl std::fmt::Display for TempId {
75    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
76        write!(f, "t{}", self.0)
77    }
78}
79
80#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
81pub struct ParamId(pub u32);
82
83impl std::fmt::Display for ParamId {
84    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
85        write!(f, "p{}", self.0)
86    }
87}
88
89#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
90pub struct BlockParamId {
91    pub block: crate::block::BlockId,
92    pub index: u32,
93}
94
95impl std::fmt::Display for BlockParamId {
96    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
97        write!(f, "{}:p{}", self.block, self.index)
98    }
99}
100
101#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
102pub struct StorageRefId(pub u32);
103
104impl std::fmt::Display for StorageRefId {
105    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
106        write!(f, "sref{}", self.0)
107    }
108}
109
110#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
111pub struct MemoryRefId(pub u32);
112
113impl std::fmt::Display for MemoryRefId {
114    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
115        write!(f, "mref{}", self.0)
116    }
117}
118
119#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
120pub struct GlobalId(pub u32);
121
122impl std::fmt::Display for GlobalId {
123    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124        write!(f, "g{}", self.0)
125    }
126}
127
128#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
129pub enum Constant {
130    Bool(bool),
131    Uint(BigUint, u16),
132    Int(BigInt, u16),
133    Address([u8; 20]),
134    Bytes(Vec<u8>),
135    String(String),
136    Null,
137}
138
139impl Constant {
140    pub fn zero(ty: &Type) -> Option<Self> {
141        match ty {
142            Type::Bool => Some(Constant::Bool(false)),
143            Type::Uint(bits) => Some(Constant::Uint(BigUint::from(0u32), *bits)),
144            Type::Int(bits) => Some(Constant::Int(BigInt::from(0), *bits)),
145            Type::Address => Some(Constant::Address([0; 20])),
146            Type::Bytes(n) => Some(Constant::Bytes(vec![0; *n as usize])),
147            _ => None,
148        }
149    }
150
151    pub fn one(ty: &Type) -> Option<Self> {
152        match ty {
153            Type::Bool => Some(Constant::Bool(true)),
154            Type::Uint(bits) => Some(Constant::Uint(BigUint::from(1u32), *bits)),
155            Type::Int(bits) => Some(Constant::Int(BigInt::from(1), *bits)),
156            _ => None,
157        }
158    }
159
160    pub fn as_int(&self) -> Option<i64> {
161        match self {
162            Constant::Uint(val, _) => val.to_u64_digits().first().copied().and_then(|v| {
163                if v <= i64::MAX as u64 {
164                    Some(v as i64)
165                } else {
166                    None
167                }
168            }),
169            Constant::Int(val, _) => val.to_i64(),
170            Constant::Bool(b) => Some(if *b { 1 } else { 0 }),
171            _ => None,
172        }
173    }
174}
175
176impl std::fmt::Display for Constant {
177    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
178        match self {
179            Constant::Bool(b) => write!(f, "{}", b),
180            Constant::Uint(val, bits) => write!(f, "{}u{}", val, bits),
181            Constant::Int(val, bits) => write!(f, "{}i{}", val, bits),
182            Constant::Address(addr) => write!(f, "0x{}", hex::encode(addr)),
183            Constant::Bytes(bytes) => write!(f, "0x{}", hex::encode(bytes)),
184            Constant::String(s) => write!(f, "\"{}\"", s),
185            Constant::Null => write!(f, "null"),
186        }
187    }
188}
189
190#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
191pub enum Location {
192    Stack { offset: i32 },
193    Memory { base: Value, offset: Value },
194    Storage { slot: Value },
195    Calldata { offset: Value },
196    ReturnData { offset: Value },
197}
198
199#[derive(Debug, Clone, Default, Serialize, Deserialize)]
200pub struct ValueMetadata {
201    pub is_tainted: bool,
202    pub is_user_input: bool,
203    pub is_constant_folded: bool,
204    pub range: Option<ValueRange>,
205    pub source_location: Option<SourceLocation>,
206}
207
208#[derive(Debug, Clone, Serialize, Deserialize)]
209pub struct ValueRange {
210    pub min: Constant,
211    pub max: Constant,
212}
213
214#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
215pub struct SourceLocation {
216    pub file: String,
217    pub line: u32,
218    pub column: u32,
219    pub end_line: Option<u32>,
220    pub end_column: Option<u32>,
221    pub start_byte: usize,
222    pub end_byte: usize,
223}
224
225impl SourceLocation {
226    pub fn new(file: String, line: u32, column: u32, start_byte: usize, end_byte: usize) -> Self {
227        Self {
228            file,
229            line,
230            column,
231            end_line: None,
232            end_column: None,
233            start_byte,
234            end_byte,
235        }
236    }
237
238    pub fn from_node(file: String, node: &tree_sitter::Node) -> Self {
239        let start = node.start_position();
240        let end = node.end_position();
241
242        Self {
243            file,
244            line: start.row as u32 + 1,
245            column: start.column as u32 + 1,
246            end_line: Some(end.row as u32 + 1),
247            end_column: Some(end.column as u32 + 1),
248            start_byte: node.start_byte(),
249            end_byte: node.end_byte(),
250        }
251    }
252
253    pub fn extract_snippet(&self, source_code: &str) -> Option<String> {
254        if self.start_byte < source_code.len() && self.end_byte <= source_code.len() {
255            let snippet = &source_code[self.start_byte..self.end_byte];
256            return Some(snippet.to_string());
257        }
258
259        let lines: Vec<&str> = source_code.lines().collect();
260        if lines.is_empty() {
261            return None;
262        }
263
264        let start_line = (self.line as usize).saturating_sub(1);
265        let end_line = self
266            .end_line
267            .map(|l| l as usize)
268            .unwrap_or(self.line as usize);
269
270        if start_line >= lines.len() {
271            return None;
272        }
273
274        let end_line = end_line.min(lines.len());
275
276        if start_line == end_line - 1 {
277            Some(lines[start_line].to_string())
278        } else {
279            Some(lines[start_line..end_line].join("\n"))
280        }
281    }
282}
283
284mod hex {
285    pub fn encode(bytes: &[u8]) -> String {
286        bytes.iter().map(|b| format!("{:02x}", b)).collect()
287    }
288}