use super::ops::{AssignmentOp, BinaryOp, UnaryOp};
use crate::address::Address;
use crate::arena::Idx;
use crate::types::TypeId;
pub type ExpressionId = Idx<ExpressionNode>;
pub type StatementId = Idx<StatementNode>;
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub enum NodeRef {
Expression(ExpressionId),
Statement(StatementId),
}
impl NodeRef {
#[inline]
#[must_use]
pub fn as_expression(self) -> Option<ExpressionId> {
match self {
NodeRef::Expression(e) => Some(e),
NodeRef::Statement(_) => None,
}
}
#[inline]
#[must_use]
pub fn as_statement(self) -> Option<StatementId> {
match self {
NodeRef::Statement(s) => Some(s),
NodeRef::Expression(_) => None,
}
}
#[inline]
#[must_use]
pub fn is_expression(self) -> bool {
matches!(self, NodeRef::Expression(_))
}
#[inline]
#[must_use]
pub fn is_statement(self) -> bool {
matches!(self, NodeRef::Statement(_))
}
}
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
#[doc(alias("lvar_t"))]
pub struct LocalId(
pub u32,
);
#[derive(Clone, PartialEq, Eq, Debug)]
#[doc(alias("argloc_t"))]
pub enum LocalLocation {
#[doc(alias("ALOC_REG1"))]
Register(u32),
#[doc(alias("ALOC_REG2"))]
RegisterPair {
low: u32,
high: u32,
},
#[doc(alias("ALOC_STACK"))]
Stack(i64),
#[doc(alias("ALOC_RREL"))]
RegisterRelative {
reg: u32,
offset: i64,
},
#[doc(alias("ALOC_STATIC"))]
Static(Address),
#[doc(alias("ALOC_DIST"))]
Scattered(Vec<LocationPiece>),
#[doc(alias("ALOC_CUSTOM"))]
Custom,
#[doc(alias("ALOC_NONE"))]
Unallocated,
}
impl LocalLocation {
pub(crate) fn from_argloc(
atype: u32,
reg1: u32,
reg2: u32,
sval: i64,
pieces: Vec<LocationPiece>,
) -> Self {
match atype {
0 => LocalLocation::Unallocated,
1 => LocalLocation::Stack(sval),
2 => LocalLocation::Scattered(pieces),
3 => LocalLocation::Register(reg1),
4 => LocalLocation::RegisterPair {
low: reg1,
high: reg2,
},
5 => LocalLocation::RegisterRelative {
reg: reg1,
offset: sval,
},
6 => LocalLocation::Static(Address::new_const(sval as u64)),
_ => LocalLocation::Custom,
}
}
}
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct LocationPiece {
pub location: LocalLocation,
pub offset: u32,
pub size: u32,
}
#[derive(Clone, Debug, PartialEq)]
#[doc(alias("lvar_t"))]
pub struct Local {
pub name: String,
pub ty: TypeId,
pub is_arg: bool,
pub is_result: bool,
pub is_byref: bool,
pub width: u32,
pub comment: Option<String>,
pub location: LocalLocation,
}
#[derive(Clone, Debug, PartialEq)]
pub struct ExpressionNode {
pub address: Option<Address>,
pub ty: TypeId,
pub parent: Option<NodeRef>,
pub kind: ExpressionKind,
}
#[derive(Clone, Debug, PartialEq)]
pub struct StatementNode {
pub address: Option<Address>,
pub parent: Option<NodeRef>,
pub kind: StatementKind,
}
#[derive(Clone, Debug, PartialEq)]
#[doc(alias("ctype_t", "cexpr_t"))]
pub enum ExpressionKind {
#[doc(alias("cot_add", "cot_sub"))]
Binary {
op: BinaryOp,
x: ExpressionId,
y: ExpressionId,
},
#[doc(alias("cot_asg"))]
Assign {
op: AssignmentOp,
x: ExpressionId,
y: ExpressionId,
},
#[doc(alias("cot_neg", "cot_lnot"))]
Unary {
op: UnaryOp,
x: ExpressionId,
},
Ternary {
cond: ExpressionId,
then_: ExpressionId,
else_: ExpressionId,
},
#[doc(alias("cot_call"))]
Call {
callee: ExpressionId,
args: Vec<ExpressionId>,
},
#[doc(alias("cot_idx"))]
Index {
array: ExpressionId,
index: ExpressionId,
},
#[doc(alias("cot_memref"))]
MemberRef {
obj: ExpressionId,
byte_offset: u32,
},
#[doc(alias("cot_memptr"))]
MemberPtr {
obj: ExpressionId,
byte_offset: u32,
},
#[doc(alias("cot_cast"))]
Cast {
x: ExpressionId,
},
#[doc(alias("cot_ptr"))]
Deref {
x: ExpressionId,
size: u32,
},
#[doc(alias("cot_sizeof"))]
Sizeof(ExpressionId),
#[doc(alias("cot_num"))]
Num(u64),
#[doc(alias("cot_fnum"))]
Fnum(f64),
#[doc(alias("cot_str"))]
Str(String),
#[doc(alias("cot_obj"))]
Obj {
address: Address,
name: Option<String>,
},
#[doc(alias("cot_var"))]
Var(LocalId),
#[doc(alias("cot_helper"))]
Helper(String),
#[doc(alias("cot_type"))]
TypeExpression,
#[doc(alias("cot_empty"))]
Empty,
Internal,
}
macro_rules! expression_accessors {
( $( $(#[$m:meta])* $fn:ident : $variant:ident $pat:tt => $ret:ty = $build:expr ; )* ) => {
impl ExpressionKind {
$(
$(#[$m])*
#[inline]
#[must_use]
pub fn $fn(&self) -> Option<$ret> {
if let ExpressionKind::$variant $pat = self {
Some($build)
} else {
None
}
}
)*
}
};
}
expression_accessors! {
as_binary: Binary { op, x, y } => (BinaryOp, ExpressionId, ExpressionId) = (*op, *x, *y);
as_assign: Assign { op, x, y } => (AssignmentOp, ExpressionId, ExpressionId) = (*op, *x, *y);
as_unary: Unary { op, x } => (UnaryOp, ExpressionId) = (*op, *x);
as_ternary: Ternary { cond, then_, else_ } => (ExpressionId, ExpressionId, ExpressionId) = (*cond, *then_, *else_);
as_call: Call { callee, args } => (ExpressionId, &[ExpressionId]) = (*callee, args.as_slice());
as_index: Index { array, index } => (ExpressionId, ExpressionId) = (*array, *index);
as_member_ref: MemberRef { obj, byte_offset } => (ExpressionId, u32) = (*obj, *byte_offset);
as_member_ptr: MemberPtr { obj, byte_offset } => (ExpressionId, u32) = (*obj, *byte_offset);
as_cast: Cast { x } => ExpressionId = *x;
as_deref: Deref { x, size } => (ExpressionId, u32) = (*x, *size);
as_sizeof: Sizeof(x) => ExpressionId = *x;
as_num: Num(v) => u64 = *v;
as_fnum: Fnum(v) => f64 = *v;
as_var: Var(v) => LocalId = *v;
as_str: Str(s) => &str = s.as_str();
as_obj: Obj { address, name } => (Address, Option<&str>) = (*address, name.as_deref());
as_helper: Helper(s) => &str = s.as_str();
}
#[derive(Clone, Debug, PartialEq)]
pub struct Case {
pub values: Vec<u64>,
pub body: StatementId,
}
#[derive(Clone, Debug, PartialEq)]
#[doc(alias("cinsn_t"))]
pub enum StatementKind {
#[doc(alias("cit_block"))]
Block(Vec<StatementId>),
#[doc(alias("cit_expr"))]
Expression(ExpressionId),
#[doc(alias("cit_if"))]
If {
cond: ExpressionId,
then_: StatementId,
else_: Option<StatementId>,
},
#[doc(alias("cit_for"))]
For {
init: Option<ExpressionId>,
cond: Option<ExpressionId>,
step: Option<ExpressionId>,
body: StatementId,
},
#[doc(alias("cit_while"))]
While {
cond: ExpressionId,
body: StatementId,
},
#[doc(alias("cit_do"))]
Do {
body: StatementId,
cond: ExpressionId,
},
#[doc(alias("cit_switch"))]
Switch {
expression: ExpressionId,
cases: Vec<Case>,
},
#[doc(alias("cit_break"))]
Break,
#[doc(alias("cit_continue"))]
Continue,
#[doc(alias("cit_return"))]
Return(Option<ExpressionId>),
#[doc(alias("cit_goto"))]
Goto {
label: i32,
},
#[doc(alias("cit_asm"))]
Asm(Vec<Address>),
#[doc(alias("cit_try"))]
Try {
body: StatementId,
catches: Vec<StatementId>,
},
#[doc(alias("cit_throw"))]
Throw(Option<ExpressionId>),
#[doc(alias("cit_empty"))]
Empty,
}
impl ExpressionKind {
pub(crate) fn for_each_child(&self, mut f: impl FnMut(NodeRef)) {
use ExpressionKind::{
Assign, Binary, Call, Cast, Deref, Index, MemberPtr, MemberRef, Sizeof, Ternary, Unary,
};
match self {
Binary { x, y, .. } | Assign { x, y, .. } => {
f(NodeRef::Expression(*x));
f(NodeRef::Expression(*y));
}
Index { array, index } => {
f(NodeRef::Expression(*array));
f(NodeRef::Expression(*index));
}
Unary { x, .. } | Cast { x } | Deref { x, .. } | Sizeof(x) => {
f(NodeRef::Expression(*x))
}
MemberRef { obj, .. } | MemberPtr { obj, .. } => f(NodeRef::Expression(*obj)),
Ternary { cond, then_, else_ } => {
f(NodeRef::Expression(*cond));
f(NodeRef::Expression(*then_));
f(NodeRef::Expression(*else_));
}
Call { callee, args } => {
f(NodeRef::Expression(*callee));
args.iter().for_each(|a| f(NodeRef::Expression(*a)));
}
Self::Num(_)
| Self::Fnum(_)
| Self::Str(_)
| Self::Obj { .. }
| Self::Var(_)
| Self::Helper(_)
| Self::TypeExpression
| Self::Empty
| Self::Internal => {}
}
}
}
impl StatementKind {
pub(crate) fn for_each_child(&self, mut f: impl FnMut(NodeRef)) {
use StatementKind::{Block, Do, Expression, For, If, Return, Switch, Throw, Try, While};
match self {
Block(statements) => statements.iter().for_each(|s| f(NodeRef::Statement(*s))),
Expression(e) => f(NodeRef::Expression(*e)),
If { cond, then_, else_ } => {
f(NodeRef::Expression(*cond));
f(NodeRef::Statement(*then_));
else_.iter().for_each(|s| f(NodeRef::Statement(*s)));
}
For {
init,
cond,
step,
body,
} => {
init.iter().for_each(|e| f(NodeRef::Expression(*e)));
cond.iter().for_each(|e| f(NodeRef::Expression(*e)));
step.iter().for_each(|e| f(NodeRef::Expression(*e)));
f(NodeRef::Statement(*body));
}
While { cond, body } => {
f(NodeRef::Expression(*cond));
f(NodeRef::Statement(*body));
}
Do { body, cond } => {
f(NodeRef::Statement(*body));
f(NodeRef::Expression(*cond));
}
Switch { expression, cases } => {
f(NodeRef::Expression(*expression));
cases.iter().for_each(|c| f(NodeRef::Statement(c.body)));
}
Return(e) | Throw(e) => e.iter().for_each(|x| f(NodeRef::Expression(*x))),
Try { body, catches } => {
f(NodeRef::Statement(*body));
catches.iter().for_each(|s| f(NodeRef::Statement(*s)));
}
Self::Break | Self::Continue | Self::Goto { .. } | Self::Asm(_) | Self::Empty => {}
}
}
}
#[cfg(test)]
mod tests {
use assert2::assert;
use super::*;
fn e(n: u32) -> ExpressionId {
Idx::from_raw(n)
}
#[test]
fn expression_accessors_project_their_variant() {
let (a, b, c) = (e(0), e(1), e(2));
assert!(
ExpressionKind::Binary {
op: BinaryOp::Add,
x: a,
y: b
}
.as_binary()
== Some((BinaryOp::Add, a, b))
);
assert!(
ExpressionKind::Assign {
op: AssignmentOp::Assign,
x: a,
y: b
}
.as_assign()
== Some((AssignmentOp::Assign, a, b))
);
assert!(
ExpressionKind::Unary {
op: UnaryOp::Neg,
x: a
}
.as_unary()
== Some((UnaryOp::Neg, a))
);
assert!(
ExpressionKind::Ternary {
cond: a,
then_: b,
else_: c
}
.as_ternary()
== Some((a, b, c))
);
assert!(ExpressionKind::Index { array: a, index: b }.as_index() == Some((a, b)));
assert!(
ExpressionKind::MemberRef {
obj: a,
byte_offset: 8
}
.as_member_ref()
== Some((a, 8))
);
assert!(
ExpressionKind::MemberPtr {
obj: a,
byte_offset: 8
}
.as_member_ptr()
== Some((a, 8))
);
assert!(ExpressionKind::Cast { x: a }.as_cast() == Some(a));
assert!(ExpressionKind::Deref { x: a, size: 4 }.as_deref() == Some((a, 4)));
assert!(ExpressionKind::Sizeof(a).as_sizeof() == Some(a));
assert!(ExpressionKind::Num(7).as_num() == Some(7));
assert!(ExpressionKind::Fnum(3.5).as_fnum() == Some(3.5));
assert!(ExpressionKind::Var(LocalId(3)).as_var() == Some(LocalId(3)));
assert!(ExpressionKind::Str("hi".into()).as_str() == Some("hi"));
assert!(ExpressionKind::Helper("h".into()).as_helper() == Some("h"));
let call = ExpressionKind::Call {
callee: a,
args: vec![b, c],
};
assert!(let Some((callee, args)) = call.as_call());
assert!(callee == a && args.len() == 2 && args[0] == b && args[1] == c);
let obj = ExpressionKind::Obj {
address: Address::new_const(0x10),
name: Some("g".into()),
};
assert!(obj.as_obj() == Some((Address::new_const(0x10), Some("g"))));
assert!(let None = ExpressionKind::Num(1).as_binary());
}
#[test]
fn node_ref_projections() {
let expression = NodeRef::Expression(e(0));
let statement = NodeRef::Statement(Idx::from_raw(0));
assert!(expression.as_expression() == Some(e(0)));
assert!(let None = expression.as_statement());
assert!(expression.is_expression() && !expression.is_statement());
assert!(statement.as_statement() == Some(Idx::from_raw(0)));
assert!(statement.is_statement() && !statement.is_expression());
}
#[test]
fn from_argloc_maps_every_atype() {
use LocalLocation::*;
assert!(LocalLocation::from_argloc(0, 0, 0, 0, vec![]) == Unallocated);
assert!(LocalLocation::from_argloc(1, 0, 0, -8, vec![]) == Stack(-8));
assert!(LocalLocation::from_argloc(3, 5, 0, 0, vec![]) == Register(5));
assert!(LocalLocation::from_argloc(4, 5, 6, 0, vec![]) == RegisterPair { low: 5, high: 6 });
assert!(
LocalLocation::from_argloc(5, 5, 0, 16, vec![])
== RegisterRelative { reg: 5, offset: 16 }
);
assert!(
LocalLocation::from_argloc(6, 0, 0, 0x1000, vec![])
== Static(Address::new_const(0x1000))
);
assert!(LocalLocation::from_argloc(7, 0, 0, 0, vec![]) == Custom);
assert!(LocalLocation::from_argloc(42, 0, 0, 0, vec![]) == Custom);
let piece = LocationPiece {
location: Stack(16),
offset: 0,
size: 8,
};
assert!(
LocalLocation::from_argloc(2, 0, 0, 0, vec![piece.clone()]) == Scattered(vec![piece])
);
}
}