use alloc::vec::Vec;
use crate::bytecode::{Chunk, ConstValue, NOMINAL_COST_MODEL, Op, OpCost, OpCostContext};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TextSize {
NotText,
Known(u32),
Unbounded,
}
impl TextSize {
pub const ZERO: TextSize = TextSize::Known(0);
pub const UNBOUNDED: TextSize = TextSize::Unbounded;
pub fn saturating_add(self, other: TextSize) -> TextSize {
match (self, other) {
(TextSize::NotText, TextSize::NotText) => TextSize::NotText,
(TextSize::NotText, t) | (t, TextSize::NotText) => t,
(TextSize::Unbounded, _) | (_, TextSize::Unbounded) => TextSize::Unbounded,
(TextSize::Known(a), TextSize::Known(b)) => match a.checked_add(b) {
Some(sum) if sum < u32::MAX => TextSize::Known(sum),
_ => TextSize::Unbounded,
},
}
}
pub fn join(self, other: TextSize) -> TextSize {
match (self, other) {
(TextSize::NotText, t) | (t, TextSize::NotText) => t,
(TextSize::Unbounded, _) | (_, TextSize::Unbounded) => TextSize::Unbounded,
(TextSize::Known(a), TextSize::Known(b)) => TextSize::Known(a.max(b)),
}
}
pub fn as_u32(self) -> u32 {
match self {
TextSize::Known(n) => n,
TextSize::Unbounded | TextSize::NotText => u32::MAX,
}
}
}
pub fn op_cost_context(lhs: TextSize, rhs: TextSize) -> OpCostContext {
OpCostContext {
lhs_text_len: lhs.as_u32(),
rhs_text_len: rhs.as_u32(),
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ChunkTextAnalysis {
pub heap_alloc: u32,
pub returns_text: bool,
pub yields_text: bool,
}
pub fn chunk_text_heap_alloc(chunk: &Chunk) -> u32 {
analyze_chunk_text(chunk, &[]).heap_alloc
}
pub fn analyze_chunk_text(chunk: &Chunk, callee_returns_text: &[bool]) -> ChunkTextAnalysis {
let mut state = TextAnalysis::new(chunk.local_count as usize);
let mut total: u32 = 0;
let mut loop_depth: u32 = 0;
let mut branch_depth: u32 = 0;
let mut returns_text = false;
let mut yields_text = false;
for op in &chunk.ops {
match op {
Op::Loop(_) => loop_depth = loop_depth.saturating_add(1),
Op::If(_) => branch_depth = branch_depth.saturating_add(1),
_ => {}
}
let conservative = loop_depth > 0 || branch_depth > 0;
match op {
Op::Return => {
let top = state.peek();
if !matches!(top, TextSize::NotText) {
returns_text = true;
}
}
Op::Yield => {
let top = state.peek();
if !matches!(top, TextSize::NotText) {
yields_text = true;
}
}
_ => {}
}
let contribution = state.apply_op(op, chunk, conservative, callee_returns_text);
total = total.saturating_add(contribution);
if total == u32::MAX {
return ChunkTextAnalysis {
heap_alloc: u32::MAX,
returns_text,
yields_text,
};
}
match op {
Op::EndLoop(_) => loop_depth = loop_depth.saturating_sub(1),
Op::EndIf => branch_depth = branch_depth.saturating_sub(1),
_ => {}
}
}
let final_top = state.peek();
if !matches!(final_top, TextSize::NotText) {
returns_text = true;
}
ChunkTextAnalysis {
heap_alloc: total,
returns_text,
yields_text,
}
}
struct TextAnalysis {
stack: Vec<TextSize>,
locals: Vec<TextSize>,
}
impl TextAnalysis {
fn new(local_count: usize) -> Self {
Self {
stack: Vec::new(),
locals: alloc::vec![TextSize::NotText; local_count],
}
}
fn pop(&mut self) -> TextSize {
self.stack.pop().unwrap_or(TextSize::NotText)
}
fn push(&mut self, size: TextSize) {
self.stack.push(size);
}
fn peek(&self) -> TextSize {
self.stack.last().copied().unwrap_or(TextSize::NotText)
}
fn apply_op(
&mut self,
op: &Op,
chunk: &Chunk,
conservative: bool,
callee_returns_text: &[bool],
) -> u32 {
let sat = |size: TextSize| -> TextSize {
if conservative {
match size {
TextSize::NotText => TextSize::NotText,
_ => TextSize::Unbounded,
}
} else {
size
}
};
match op {
Op::Const(idx) => {
let size = match chunk.constants.get(*idx as usize) {
Some(ConstValue::StaticStr(s)) => {
let len = u32::try_from(s.len()).unwrap_or(u32::MAX - 1);
TextSize::Known(len.min(u32::MAX - 1))
}
Some(_) => TextSize::NotText,
None => TextSize::NotText,
};
self.push(sat(size));
0
}
Op::Add => {
let b = self.pop();
let a = self.pop();
let (result, dynamic_cost) =
if matches!(a, TextSize::NotText) || matches!(b, TextSize::NotText) {
(TextSize::NotText, 0)
} else {
let context = op_cost_context(a, b);
let cost = match NOMINAL_COST_MODEL.heap_alloc_cost(op, chunk) {
OpCost::Dynamic(f) => f(&context),
OpCost::Fixed(n) => n,
};
(a.saturating_add(b), cost)
};
self.push(sat(result));
dynamic_cost
}
Op::GetLocal(i) => {
let size = self
.locals
.get(*i as usize)
.copied()
.unwrap_or(TextSize::NotText);
self.push(sat(size));
0
}
Op::SetLocal(i) => {
let value = self.pop();
if let Some(slot) = self.locals.get_mut(*i as usize) {
*slot = sat(value);
}
0
}
Op::PopN(n) => {
for _ in 0..*n {
self.pop();
}
0
}
Op::Dup => {
let top = self.peek();
self.push(top);
0
}
Op::Call(callee_idx, n_args) => {
for _ in 0..*n_args {
self.pop();
}
let returns_text = callee_returns_text
.get(*callee_idx as usize)
.copied()
.unwrap_or(true);
let pushed = if returns_text {
TextSize::Unbounded
} else {
TextSize::NotText
};
self.push(sat(pushed));
0
}
Op::CallVerifiedNative(_, n_args) | Op::CallExternalNative(_, n_args) => {
for _ in 0..*n_args {
self.pop();
}
self.push(TextSize::Unbounded);
0
}
Op::Yield => {
self.pop();
0
}
Op::Return => {
self.pop();
0
}
Op::If(_) => {
self.pop();
0
}
Op::Else(_) | Op::EndIf | Op::Loop(_) | Op::EndLoop(_) => 0,
Op::Break(_) => 0,
Op::BreakIf(_) => {
self.pop();
0
}
_ => {
let shrink = op.stack_shrink() as usize;
for _ in 0..shrink {
self.pop();
}
let growth = op.stack_growth() as usize;
for _ in 0..growth {
self.push(TextSize::NotText);
}
0
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn add_known_values() {
assert_eq!(
TextSize::Known(3).saturating_add(TextSize::Known(5)),
TextSize::Known(8)
);
assert_eq!(
TextSize::ZERO.saturating_add(TextSize::ZERO),
TextSize::ZERO
);
}
#[test]
fn add_saturates_to_unbounded_on_overflow() {
assert_eq!(
TextSize::Known(u32::MAX - 1).saturating_add(TextSize::Known(1)),
TextSize::Unbounded
);
assert_eq!(
TextSize::Known(u32::MAX / 2 + 1).saturating_add(TextSize::Known(u32::MAX / 2 + 1)),
TextSize::Unbounded
);
}
#[test]
fn add_propagates_unbounded() {
assert_eq!(
TextSize::Unbounded.saturating_add(TextSize::Known(5)),
TextSize::Unbounded
);
assert_eq!(
TextSize::Known(5).saturating_add(TextSize::Unbounded),
TextSize::Unbounded
);
assert_eq!(
TextSize::Unbounded.saturating_add(TextSize::Unbounded),
TextSize::Unbounded
);
}
#[test]
fn join_known_takes_max() {
assert_eq!(
TextSize::Known(3).join(TextSize::Known(5)),
TextSize::Known(5)
);
assert_eq!(
TextSize::Known(7).join(TextSize::Known(2)),
TextSize::Known(7)
);
}
#[test]
fn join_propagates_unbounded() {
assert_eq!(
TextSize::Unbounded.join(TextSize::Known(5)),
TextSize::Unbounded
);
assert_eq!(
TextSize::Known(5).join(TextSize::Unbounded),
TextSize::Unbounded
);
}
#[test]
fn as_u32_projects_unbounded_to_max() {
assert_eq!(TextSize::Known(42).as_u32(), 42);
assert_eq!(TextSize::Unbounded.as_u32(), u32::MAX);
assert_eq!(TextSize::ZERO.as_u32(), 0);
}
#[test]
fn op_cost_context_carries_lattice_values() {
let ctx = op_cost_context(TextSize::Known(10), TextSize::Known(20));
assert_eq!(ctx.lhs_text_len, 10);
assert_eq!(ctx.rhs_text_len, 20);
let ctx_unbounded = op_cost_context(TextSize::Unbounded, TextSize::Known(5));
assert_eq!(ctx_unbounded.lhs_text_len, u32::MAX);
assert_eq!(ctx_unbounded.rhs_text_len, 5);
}
use crate::bytecode::{BlockType, Chunk, ConstValue, Op};
use alloc::string::String as StdString;
use alloc::vec;
fn make_chunk(ops: Vec<Op>, constants: Vec<ConstValue>, locals: u16) -> Chunk {
Chunk {
name: StdString::from("test"),
ops,
constants,
struct_templates: Vec::new(),
local_count: locals,
param_count: 0,
block_type: BlockType::Func,
param_types: Vec::new(),
}
}
#[test]
fn text_heap_alloc_empty_chunk_is_zero() {
let chunk = make_chunk(vec![Op::Return], vec![], 0);
assert_eq!(chunk_text_heap_alloc(&chunk), 0);
}
#[test]
fn text_heap_alloc_static_literal_alone_is_zero() {
let chunk = make_chunk(
vec![Op::Const(0), Op::Return],
vec![ConstValue::StaticStr(StdString::from("hello"))],
0,
);
assert_eq!(chunk_text_heap_alloc(&chunk), 0);
}
#[test]
fn text_heap_alloc_single_concat_returns_sum_of_lengths() {
let chunk = make_chunk(
vec![Op::Const(0), Op::Const(1), Op::Add, Op::Return],
vec![
ConstValue::StaticStr(StdString::from("ab")),
ConstValue::StaticStr(StdString::from("cdef")),
],
0,
);
assert_eq!(chunk_text_heap_alloc(&chunk), 6);
}
#[test]
fn text_heap_alloc_doubling_pattern_accumulates_geometric_series() {
let chunk = make_chunk(
vec![
Op::Const(0), Op::SetLocal(0), Op::GetLocal(0),
Op::GetLocal(0),
Op::Add, Op::SetLocal(0), Op::GetLocal(0),
Op::GetLocal(0),
Op::Add, Op::SetLocal(0), Op::Return,
],
vec![ConstValue::StaticStr(StdString::from("a"))],
1,
);
assert_eq!(chunk_text_heap_alloc(&chunk), 6);
}
#[test]
fn text_heap_alloc_saturates_for_long_doubling_chain() {
let mut ops = vec![Op::Const(0), Op::SetLocal(0)];
for _ in 0..60 {
ops.push(Op::GetLocal(0));
ops.push(Op::GetLocal(0));
ops.push(Op::Add);
ops.push(Op::SetLocal(0));
}
ops.push(Op::Return);
let chunk = make_chunk(ops, vec![ConstValue::StaticStr(StdString::from("a"))], 1);
assert_eq!(chunk_text_heap_alloc(&chunk), u32::MAX);
}
#[test]
fn text_heap_alloc_inside_loop_saturates_to_unbounded() {
let chunk = make_chunk(
vec![
Op::Loop(6), Op::Const(0),
Op::Const(0),
Op::Add,
Op::PopN(1),
Op::EndLoop(0), Op::Return,
],
vec![ConstValue::StaticStr(StdString::from("a"))],
0,
);
assert_eq!(chunk_text_heap_alloc(&chunk), u32::MAX);
}
#[test]
fn text_heap_alloc_call_native_result_is_unbounded() {
let chunk = make_chunk(
vec![
Op::Const(0),
Op::CallVerifiedNative(0, 1), Op::Const(0),
Op::Add, Op::Return,
],
vec![ConstValue::StaticStr(StdString::from("a"))],
0,
);
assert_eq!(chunk_text_heap_alloc(&chunk), u32::MAX);
}
#[test]
fn doubling_pattern_saturates_after_thirty_two_iterations() {
let mut size = TextSize::Known(1);
for _ in 0..31 {
size = size.saturating_add(size);
}
assert_eq!(size, TextSize::Known(1u32 << 31));
size = size.saturating_add(size);
assert_eq!(size, TextSize::Unbounded);
for _ in 0..30 {
size = size.saturating_add(size);
assert_eq!(size, TextSize::Unbounded);
}
}
}