use std::{
cell::RefCell,
collections::HashMap,
marker::PhantomData,
sync::{Arc, Weak},
};
use crate::compiler::{
compiled_code::Symbol,
graph::{node::Literal, OpCode},
};
use super::{
graph::Graph,
node::NodeId,
op::BaseOp,
ty::{AnyTy, IntTy, NumTy, Top, Ty, Type},
FloatTy, InputKind,
};
pub struct BaseGraphBuilder {
pub graph: Box<Graph>,
vars: Vec<Weak<Variable>>,
control: Option<NodeId<()>>,
effect: Option<NodeId>,
}
impl BaseGraphBuilder {
pub fn new(graph: Box<Graph>) -> Self {
Self {
graph,
vars: vec![],
control: None,
effect: None,
}
}
pub fn finalize(self) -> Box<Graph> {
trace!(target: "graph", "(original)\n{:?}", self.graph);
self.graph
}
pub fn control(&self) -> NodeId<()> {
self.control.unwrap()
}
pub fn effect(&self) -> NodeId {
self.effect.unwrap()
}
pub fn set_control(&mut self, n: Option<NodeId<()>>) {
self.control = n;
}
pub fn set_effect(&mut self, n: NodeId) {
self.effect = Some(n);
}
fn untyped_node(
&mut self,
t: Type,
op: BaseOp,
ctrl: &[NodeId<()>],
effect: &[NodeId],
inputs: &[NodeId],
) -> NodeId<Top> {
self.graph.new_untyped_node(t, op, inputs, ctrl, effect)
}
fn node<T: Ty>(
&mut self,
op: BaseOp,
ctrl: &[NodeId<()>],
effect: &[NodeId],
inputs: &[NodeId],
) -> NodeId<T> {
self.untyped_node(T::TYPE, op, ctrl, effect, inputs).cast()
}
pub fn new_node<T: Ty>(&mut self, op: BaseOp) -> NodeId<T> {
self.node(op, &[], &[], &[])
}
pub fn new_start(&mut self) -> NodeId<()> {
let start = self.node(BaseOp::Start, &[], &[], &[]);
self.set_effect(start.cast());
start
}
pub fn new_param<T: AnyTy>(&mut self, index: usize) -> NodeId<T> {
debug_assert!(T::TYPE == Type::Top || self.graph.signature.0[index] == T::TYPE);
let n = self.untyped_node(self.graph.signature.0[index], BaseOp::Param, &[], &[], &[]);
self.graph[n].literal = Some(Literal::ParamIndex(index));
n.cast()
}
pub fn new_region(&mut self, controls: &[NodeId<()>]) -> NodeId<()> {
self.node(BaseOp::Region, controls, &[], &[])
}
pub fn new_effect_phi(&mut self, effects: &[NodeId]) -> NodeId<()> {
let ctrl = self.control();
let e = self.node(BaseOp::EffectPhi, &[ctrl], effects, &[]);
self.set_effect(e.cast());
e
}
pub fn new_phi<T: Ty>(&mut self) -> NodeId<T> {
self.node(BaseOp::Phi, &[], &[], &[])
}
pub fn new_phi_with_type(&mut self, t: Type) -> NodeId {
self.untyped_node(t, BaseOp::Phi, &[], &[], &[])
}
pub fn new_constant<T: NumTy>(&mut self, value: T) -> NodeId<T> {
let n = self.node(BaseOp::Const, &[], &[], &[]);
self.graph[n].literal = Some(Literal::Value(T::constant(value)));
n
}
pub fn new_not<T: IntTy>(&mut self, x: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::Not, &[], &[], &[x.cast()])
}
pub fn new_or<T: IntTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::Or, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_xor<T: IntTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::Xor, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_and<T: IntTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::And, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_itrunc_u<T: IntTy, U: IntTy>(&mut self, x: NodeId<T>) -> NodeId<U> {
assert!(T::TYPE.mem_size() > U::TYPE.mem_size());
self.node(BaseOp::ITruncU, &[], &[], &[x.cast()])
}
pub fn new_cvt_f2si<T: FloatTy, U: IntTy>(&mut self, x: NodeId<T>) -> NodeId<U> {
self.node(BaseOp::CvtF2SI, &[], &[], &[x.cast()])
}
pub fn new_cvt_f2ui<T: FloatTy, U: IntTy>(&mut self, x: NodeId<T>) -> NodeId<U> {
self.node(BaseOp::CvtF2UI, &[], &[], &[x.cast()])
}
pub fn new_shl<T: NumTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::Shl, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_shr_s<T: NumTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::ShrS, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_shr_u<T: NumTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::ShrU, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_rotl<T: NumTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::Rotl, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_rotr<T: NumTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::Rotr, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_add<T: NumTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::Add, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_sub<T: NumTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::Sub, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_mul<T: NumTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::Mul, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_div_s<T: IntTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<T> {
assert_ne!(T::TYPE, Type::Bool);
self.node(BaseOp::DivS, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_div_u<T: IntTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<T> {
assert_ne!(T::TYPE, Type::Bool);
self.node(BaseOp::DivU, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_rem_s<T: IntTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::RemS, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_rem_u<T: IntTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::RemU, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_clz<T: IntTy>(&mut self, x: NodeId<T>) -> NodeId<T> {
assert_ne!(T::TYPE, Type::Bool);
assert_ne!(T::TYPE, Type::I8);
self.node(BaseOp::Clz, &[], &[], &[x.cast()])
}
pub fn new_ctz<T: IntTy>(&mut self, x: NodeId<T>) -> NodeId<T> {
assert_ne!(T::TYPE, Type::Bool);
assert_ne!(T::TYPE, Type::I8);
self.node(BaseOp::Ctz, &[], &[], &[x.cast()])
}
pub fn new_popcnt<T: IntTy>(&mut self, x: NodeId<T>) -> NodeId<T> {
assert_ne!(T::TYPE, Type::Bool);
assert_ne!(T::TYPE, Type::I8);
self.node(BaseOp::Popcnt, &[], &[], &[x.cast()])
}
pub fn new_fdiv<T: FloatTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::FDiv, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_fsqrt<T: FloatTy>(&mut self, x: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::FSqrt, &[], &[], &[x.cast()])
}
pub fn new_fround<T: FloatTy>(&mut self, x: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::FRound, &[], &[], &[x.cast()])
}
pub fn new_ffloor<T: FloatTy>(&mut self, x: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::FFloor, &[], &[], &[x.cast()])
}
pub fn new_fceil<T: FloatTy>(&mut self, x: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::FCeil, &[], &[], &[x.cast()])
}
pub fn new_ftrunc<T: FloatTy>(&mut self, x: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::FTrunc, &[], &[], &[x.cast()])
}
pub fn new_fabs<T: FloatTy>(&mut self, x: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::FAbs, &[], &[], &[x.cast()])
}
pub fn new_fcopysign<T: FloatTy>(&mut self, mag: NodeId<T>, sgn: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::FCopysign, &[], &[], &[mag.cast(), sgn.cast()])
}
pub fn new_neg<T: NumTy>(&mut self, x: NodeId<T>) -> NodeId<T> {
self.node(BaseOp::Neg, &[], &[], &[x.cast()])
}
pub fn new_eq<T: NumTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<bool> {
self.node(BaseOp::Eq, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_ne<T: NumTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<bool> {
self.node(BaseOp::Ne, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_lt_s<T: IntTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<bool> {
self.node(BaseOp::LtS, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_lt_u<T: IntTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<bool> {
self.node(BaseOp::LtU, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_le_s<T: IntTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<bool> {
self.node(BaseOp::LeS, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_le_u<T: IntTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<bool> {
self.node(BaseOp::LeU, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_gt_s<T: IntTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<bool> {
self.node(BaseOp::GtS, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_gt_u<T: IntTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<bool> {
self.node(BaseOp::GtU, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_ge_s<T: IntTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<bool> {
self.node(BaseOp::GeS, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_ge_u<T: IntTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<bool> {
self.node(BaseOp::GeU, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_zext<T: IntTy, U: IntTy>(&mut self, x: NodeId<T>) -> NodeId<U> {
self.node(BaseOp::ZExt, &[], &[], &[x.cast()])
}
pub fn new_sext<T: IntTy, U: IntTy>(&mut self, x: NodeId<T>) -> NodeId<U> {
self.node(BaseOp::SExt, &[], &[], &[x.cast()])
}
pub fn new_is_nan<T: FloatTy>(&mut self, x: NodeId<T>) -> NodeId<bool> {
self.node(BaseOp::IsNan, &[], &[], &[x.cast()])
}
pub fn new_lt_f<T: FloatTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<bool> {
self.node(BaseOp::LtF, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_le_f<T: FloatTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<bool> {
self.node(BaseOp::LeF, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_gt_f<T: FloatTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<bool> {
self.node(BaseOp::GtF, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_ge_f<T: FloatTy>(&mut self, x: NodeId<T>, y: NodeId<T>) -> NodeId<bool> {
self.node(BaseOp::GeF, &[], &[], &[x.cast(), y.cast()])
}
pub fn new_cvt_si2f<T: IntTy, U: FloatTy>(&mut self, x: NodeId<T>) -> NodeId<U> {
self.node(BaseOp::CvtSI2F, &[], &[], &[x.cast()])
}
pub fn new_cvt_ui2f<T: IntTy, U: FloatTy>(&mut self, x: NodeId<T>) -> NodeId<U> {
self.node(BaseOp::CvtUI2F, &[], &[], &[x.cast()])
}
pub fn new_cvt_f2f<T: FloatTy, U: FloatTy>(&mut self, x: NodeId<T>) -> NodeId<U> {
self.node(BaseOp::CvtF2F, &[], &[], &[x.cast()])
}
pub fn new_wrap<T: IntTy, U: IntTy>(&mut self, x: NodeId<T>) -> NodeId<U> {
self.node(BaseOp::Wrap, &[], &[], &[x.cast()])
}
pub fn new_bitcast<T: NumTy, U: NumTy>(&mut self, x: NodeId<T>) -> NodeId<U> {
assert_eq!(T::TYPE.mem_size(), U::TYPE.mem_size());
self.node(BaseOp::Bitcast, &[], &[], &[x.cast()])
}
pub fn new_debug_break(&mut self) -> NodeId<()> {
let ctrl = self.control();
let effect = self.effect();
let n = self.node(BaseOp::DebugBreak, &[ctrl], &[effect], &[]);
self.set_effect(n.cast());
n
}
pub fn new_return<T: AnyTy>(&mut self, v: Option<NodeId<T>>) -> NodeId<()> {
let ctrl = self.control();
self.node(
BaseOp::Return,
&[ctrl],
&[],
&v.map(|x| vec![x.cast()]).unwrap_or_else(|| vec![]),
)
}
pub fn new_call<T: Ty>(&mut self, symbol: Symbol, args: &[NodeId]) -> NodeId<T> {
let ctrl: NodeId<()> = self.control();
let effect = self.effect();
let n = self.node(BaseOp::Call, &[ctrl], &[effect], args);
self.graph[n].literal = Some(Literal::Func(symbol));
self.set_effect(n.cast());
n
}
pub fn new_call_indirect<T: Ty>(
&mut self,
fptr: NodeId<i64>,
ctx: Option<NodeId<i64>>,
args: &[NodeId],
) -> NodeId<T> {
let ctrl = self.control();
let effect = self.effect();
let mut inputs = vec![fptr.cast::<Top>()];
if let Some(ctx) = ctx {
inputs.push(ctx.cast());
}
inputs.append(&mut args.to_vec());
let n = self.node(BaseOp::CallIndirect, &[ctrl], &[effect], &inputs);
self.set_effect(n.cast());
self.graph[n].has_call_indirect_ctx = ctx.is_some();
n
}
pub fn new_load<T: Ty>(&mut self, pointer: NodeId<i64>) -> NodeId<T> {
let ctrl = self.control();
let effect = self.effect();
let n = self.node(BaseOp::Load, &[ctrl], &[effect], &[pointer.cast()]);
self.set_effect(n.cast());
n
}
pub fn new_store<T: Ty>(&mut self, pointer: NodeId<i64>, value: NodeId<T>) -> NodeId<()> {
debug_assert_ne!(T::TYPE, Type::Void);
let ctrl = self.control();
let effect = self.effect();
let n = self.node(
BaseOp::Store,
&[ctrl],
&[effect],
&[pointer.cast(), value.cast()],
);
self.set_effect(n.cast());
n
}
pub fn new_variable(&mut self, initial: NodeId) -> Arc<Variable> {
let ty = self.graph[initial].ty;
let v = Variable::new(self, ty);
v.set(self.control(), initial);
v
}
fn assert_teriminal_node_is_not_set(&self) {
let ctrl = self.control();
assert!(
self.graph[ctrl]
.control_uses
.iter()
.all(|x| self.graph[x.user()]
.op::<BaseOp>()
.ctrl_op()
.map(|o| !o.is_terminal())
.unwrap_or(true)),
"Already has a terminal node: {:?}",
self.graph[ctrl]
.control_uses
.iter()
.map(|u| u.user())
.collect::<Vec<_>>(),
);
}
pub fn new_jump(&mut self, label: NodeId<()>) -> NodeId<()> {
self.assert_teriminal_node_is_not_set();
let ctrl = self.control();
let n = self.node(BaseOp::Jump, &[ctrl], &[], &[]);
self.graph
.add_input(label.cast(), InputKind::Control, n.cast());
n
}
pub fn new_branch(
&mut self,
cond: NodeId<bool>,
then_label: NodeId<()>,
else_label: NodeId<()>,
) -> NodeId<()> {
self.assert_teriminal_node_is_not_set();
let ctrl = self.control();
let n = self.node(BaseOp::Branch, &[ctrl], &[], &[cond.cast()]);
self.graph
.add_input(then_label.cast(), InputKind::Control, n.cast());
self.graph
.add_input(else_label.cast(), InputKind::Control, n.cast());
n
}
pub fn new_br_table(
&mut self,
index: NodeId<i32>,
targets: &[NodeId<()>],
default: NodeId<()>,
) -> NodeId<()> {
self.assert_teriminal_node_is_not_set();
let ctrl = self.control();
let n = self.node(BaseOp::BrTable, &[ctrl], &[], &[index.cast()]);
for t in targets {
self.graph.add_input(t.cast(), InputKind::Control, n.cast());
}
self.graph
.add_input(default.cast(), InputKind::Control, n.cast());
n
}
pub fn r#if(&mut self, cond: NodeId<bool>) -> BranchBuilder<'_> {
BranchBuilder::new(self, cond)
}
pub fn r#loop(&mut self) -> LoopBuilder<'_> {
LoopBuilder::new(self)
}
pub fn bind(&mut self, region: NodeId<()>) {
self.set_control(Some(region));
if self.graph[region].op::<BaseOp>() == BaseOp::Start || !self.graph[region].never_binded {
return;
}
self.graph[region].never_binded = false;
let preds = self.graph[region]
.controls
.iter()
.map(|x| self.graph[*x].controls[0])
.collect::<Vec<_>>();
for v in self.vars.to_vec() {
if let Some(v) = v.upgrade() {
let phi = self.new_phi_with_type(v.ty);
self.graph
.add_input(phi.cast(), InputKind::Control, region.cast());
for pred in &preds {
let vt = self.graph[v.values.borrow()[pred]].ty;
assert_eq!(v.ty, vt);
self.graph.add_input(
phi.cast(),
InputKind::Data,
v.values.borrow()[pred].cast(),
);
}
v.set(self.control(), phi);
v.phis.borrow_mut().insert(region, phi);
}
}
}
pub fn loop_back(&mut self, loop_start: NodeId<()>) -> NodeId<()> {
let ctrl = self.control();
let n = self.node(BaseOp::Jump, &[ctrl], &[], &[]);
self.graph
.add_input(loop_start.cast(), InputKind::Control, n.cast());
for v in self.vars.to_vec() {
if let Some(v) = v.upgrade() {
let current_value = v.get();
if let Some(phi) = v.phis.borrow().get(&loop_start).cloned() {
self.graph
.add_input(phi.cast(), InputKind::Data, current_value.cast());
}
}
}
n
}
}
pub struct BranchBuilder<'a> {
builder: &'a mut BaseGraphBuilder,
br: NodeId<()>,
final_block: NodeId<()>,
has_else: bool,
effects: Vec<NodeId>,
phantom: PhantomData<&'a ()>,
}
impl<'a> BranchBuilder<'a> {
pub fn new(builder: &'a mut BaseGraphBuilder, cond: NodeId<bool>) -> Self {
let ctrl = builder.control();
let br = builder.node(BaseOp::Branch, &[ctrl], &[], &[cond.cast()]);
let final_block = builder.new_region(&[]);
let effect = builder.effect();
Self {
builder,
br,
final_block,
has_else: false,
effects: vec![effect],
phantom: PhantomData,
}
}
pub fn then_(mut self, mut f: impl FnMut(&mut BaseGraphBuilder)) -> Self {
let r = self.builder.new_region(&[self.br]);
self.builder.bind(r);
f(self.builder);
self.builder.new_jump(self.final_block);
self.effects.push(self.builder.effect());
self
}
pub fn else_(mut self, mut f: impl FnMut(&mut BaseGraphBuilder)) -> Self {
self.has_else = true;
let r = self.builder.new_region(&[self.br]);
self.builder.bind(r);
f(self.builder);
self.builder.new_jump(self.final_block);
self.effects.push(self.builder.effect());
self
}
pub fn finish(self) {
if !self.has_else {
self.builder.graph.add_input(
self.final_block.cast(),
InputKind::Control,
self.br.cast(),
);
}
self.builder.bind(self.final_block);
if self.effects.len() == 2 {
self.builder
.new_effect_phi(&[self.effects[1], self.effects[0]]);
} else {
self.builder
.new_effect_phi(&[self.effects[1], self.effects[2]]);
}
}
}
pub struct LoopBuilder<'a> {
builder: &'a mut BaseGraphBuilder,
}
impl<'a> LoopBuilder<'a> {
pub fn new(builder: &'a mut BaseGraphBuilder) -> Self {
Self { builder }
}
pub fn body(self, mut f: impl FnMut(&mut BaseGraphBuilder)) {
let loop_start = self.builder.new_region(&[]);
self.builder.new_jump(loop_start);
self.builder.bind(loop_start);
let initial_effect = self.builder.effect();
let effect_phi = self.builder.new_effect_phi(&[initial_effect]);
f(self.builder);
self.builder.loop_back(loop_start);
let next_effect = self.builder.effect();
self.builder
.graph
.add_input(effect_phi.cast(), InputKind::Effect, next_effect);
}
}
pub struct Variable {
pub ty: Type,
pub phis: RefCell<HashMap<NodeId<()>, NodeId>>,
values: RefCell<HashMap<NodeId<()>, NodeId>>,
value: RefCell<Option<NodeId>>,
}
impl Variable {
fn new(builder: &mut BaseGraphBuilder, ty: Type) -> Arc<Self> {
debug_assert_ne!(ty, Top::TYPE);
let var = Arc::new(Self {
ty,
value: RefCell::new(None),
phis: RefCell::new(HashMap::new()),
values: RefCell::new(HashMap::new()),
});
builder.vars.push(Arc::downgrade(&var));
var
}
pub fn get(&self) -> NodeId {
self.value.borrow().unwrap()
}
pub fn set(&self, ctrl: NodeId<()>, v: NodeId) {
*self.value.borrow_mut() = Some(v);
self.values.borrow_mut().insert(ctrl, v);
}
}