use alloc::{collections::BTreeMap, format, string::String, vec::Vec};
use iridia_ir::{Block, Body, End, Function, LabelId, Node, Region, Target, Terminator, ValueId};
#[derive(Clone, Copy)]
struct Exits {
exit: u32,
header: Option<u32>,
}
struct Flat {
blocks: Vec<Block>,
labels: BTreeMap<LabelId, Exits>,
at: usize,
}
impl Flat {
fn block(&mut self, params: Vec<ValueId>) -> Result<u32, String> {
let id = u32::try_from(self.blocks.len()).map_err(|_| String::from("more than 4294967295 blocks"))?;
self.blocks.push(Block {
params,
nodes: Vec::new(),
term: Terminator::Unreachable,
});
Ok(id)
}
fn current(&mut self) -> Result<&mut Block, String> {
self.blocks
.get_mut(self.at)
.ok_or_else(|| String::from("the block being filled is missing"))
}
fn end(&mut self, term: Terminator, next: u32) -> Result<(), String> {
self.current()?.term = term;
self.at = next as usize;
Ok(())
}
fn exits(&self, label: LabelId) -> Result<Exits, String> {
self.labels
.get(&label)
.copied()
.ok_or_else(|| format!("an exit to label {}, which no construct around it has", label.0))
}
fn region(&mut self, r: Region, join: Option<u32>) -> Result<(), String> {
for node in r.nodes {
match node {
Node::Value(_) | Node::StoreVar(..) | Node::StoreOutput(..) | Node::Demote => {
self.current()?.nodes.push(node)
}
Node::If {
cond,
accept,
reject,
results,
} => {
let after = self.block(results)?;
let (a, b) = (self.block(Vec::new())?, self.block(Vec::new())?);
let to = |block| Target {
block,
args: Vec::new(),
};
self.end(
Terminator::Branch {
cond,
accept: to(a),
reject: to(b),
},
a,
)?;
self.region(accept, Some(after))?;
self.at = b as usize;
self.region(reject, Some(after))?;
self.at = after as usize;
}
Node::Block { label, body, results } => {
let exit = self.block(results)?;
self.labels.insert(label, Exits { exit, header: None });
self.region(body, Some(exit))?;
self.at = exit as usize;
}
Node::Loop {
label,
params,
inits,
body,
results,
} => {
let exit = self.block(results)?;
let header = self.block(params)?;
self.labels.insert(
label,
Exits {
exit,
header: Some(header),
},
);
self.end(
Terminator::Jump(Target {
block: header,
args: inits,
}),
header,
)?;
self.region(body, Some(exit))?;
self.at = exit as usize;
}
}
}
let term = match r.end {
Some(End::Yield(args)) => match join {
Some(block) => Terminator::Jump(Target { block, args }),
None => Terminator::Return,
},
Some(End::Br(label, args)) => Terminator::Jump(Target {
block: self.exits(label)?.exit,
args,
}),
Some(End::Continue(label, args)) => Terminator::Jump(Target {
block: self
.exits(label)?
.header
.ok_or_else(|| format!("a continue of label {}, which is not a loop's", label.0))?,
args,
}),
Some(End::Return) => Terminator::Return,
Some(End::Kill) => Terminator::Kill,
Some(End::Unreachable) => Terminator::Unreachable,
None => return Err("a region with no end".into()),
};
let rest = self.block(Vec::new())?;
self.end(term, rest)
}
}
pub fn flatten(f: &mut Function) -> Result<(), String> {
let r = match core::mem::replace(&mut f.body, Body::Cfg(Vec::new())) {
Body::Regions(r) => r,
graph => {
f.body = graph;
return Ok(());
}
};
let mut flat = Flat {
blocks: Vec::new(),
labels: BTreeMap::new(),
at: 0,
};
flat.block(Vec::new())?;
flat.region(r, None)?;
f.body = Body::Cfg(flat.blocks);
Ok(())
}