use crate::sampling::alias::AliasTableU64;
use polydat::ast::CompiledU64Op;
use polydat::compile::fusion::{DecomposedGraph, DecomposedWire};
use polydat::derive_support::Config;
fn parse_weighted_str_spec(spec: &str) -> (Vec<String>, Vec<f64>) {
let mut values = Vec::new();
let mut weights = Vec::new();
for elem in spec.split([';', ',']) {
let elem = elem.trim();
if elem.is_empty() {
continue;
}
let parts: Vec<&str> = elem.splitn(2, ':').collect();
assert_eq!(parts.len(), 2, "expected 'value:weight', got '{elem}'");
values.push(parts[0].to_string());
weights.push(parts[1].parse::<f64>().expect("invalid weight"));
}
(values, weights)
}
fn parse_weighted_u64_spec(spec: &str) -> (Vec<u64>, Vec<f64>) {
let mut values = Vec::new();
let mut weights = Vec::new();
for elem in spec.split([';', ',']) {
let elem = elem.trim();
if elem.is_empty() {
continue;
}
let parts: Vec<&str> = elem.splitn(2, ':').collect();
assert_eq!(parts.len(), 2, "expected 'value:weight', got '{elem}'");
values.push(parts[0].parse::<u64>().expect("invalid value"));
weights.push(parts[1].parse::<f64>().expect("invalid weight"));
}
(values, weights)
}
pub struct WeightedStrCache {
values: Vec<String>,
table: AliasTableU64,
}
impl polydat::derive_support::PolydatSetup for WeightedStrCache {}
fn build_weighted_str_cache(spec: &str) -> WeightedStrCache {
let (values, weights) = parse_weighted_str_spec(spec);
let table = AliasTableU64::from_weights(&weights);
WeightedStrCache { values, table }
}
#[polydat::polydat_node(category = Weighted)]
fn weighted_strings(
input: u64,
spec: polydat::derive_support::Const<&str>,
#[poly_const(build_weighted_str_cache, from = spec)] cache: &WeightedStrCache,
) -> String {
let _ = spec; let idx = cache.table.sample(input) as usize;
cache.values[idx].clone()
}
pub struct WeightedU64Cache {
values: Vec<u64>,
table: AliasTableU64,
}
impl polydat::derive_support::PolydatSetup for WeightedU64Cache {}
fn build_weighted_u64_cache(spec: &str) -> WeightedU64Cache {
let (values, weights) = parse_weighted_u64_spec(spec);
let table = AliasTableU64::from_weights(&weights);
WeightedU64Cache { values, table }
}
#[polydat::polydat_node(category = Weighted)]
fn weighted_u64(
input: u64,
spec: polydat::derive_support::Const<&str>,
#[poly_const(build_weighted_u64_cache, from = spec)] cache: &WeightedU64Cache,
) -> u64 {
let _ = spec; let idx = cache.table.sample(input) as usize;
cache.values[idx]
}
pub struct WeightedPickState {
pub table: AliasTableU64,
pub values: Vec<u64>,
pub weights: Vec<f64>,
}
impl polydat::derive_support::PolydatSetup for WeightedPickState {}
fn parse_weighted_pick_spec(spec: &str) -> WeightedPickState {
let mut weights = Vec::new();
let mut values = Vec::new();
for entry in spec.split([';', ',']) {
let entry = entry.trim();
if entry.is_empty() {
continue;
}
let (v, w) = entry.split_once(':').unwrap_or_else(|| {
panic!("weighted_pick: malformed entry '{entry}', expected 'value:weight'")
});
let value: u64 = v
.trim()
.parse()
.unwrap_or_else(|_| panic!("weighted_pick: invalid value '{v}' in entry '{entry}'"));
let weight: f64 = w
.trim()
.parse()
.unwrap_or_else(|_| panic!("weighted_pick: invalid weight '{w}' in entry '{entry}'"));
assert!(
weight.is_finite() && weight > 0.0,
"weighted_pick: weight must be a positive finite f64, got {weight}",
);
values.push(value);
weights.push(weight);
}
assert!(
!weights.is_empty(),
"weighted_pick requires at least one entry in spec",
);
WeightedPickState {
table: AliasTableU64::from_weights(&weights),
values,
weights,
}
}
fn weighted_pick_jit(node: &WeightedPick) -> CompiledU64Op {
let values = node.state.values.clone();
let biases = node.state.table.biases().to_vec();
let primaries = node.state.table.primaries().to_vec();
let aliases = node.state.table.aliases().to_vec();
let n = values.len();
Box::new(move |inputs, outputs| {
let input = inputs[0];
let slot = (input as usize) % n;
let bias_test = ((input >> 32) as f64) / (u32::MAX as f64);
let index = if bias_test < biases[slot] {
primaries[slot]
} else {
aliases[slot]
};
outputs[0] = values[index as usize];
})
}
fn weighted_pick_jit_constants(node: &WeightedPick) -> Vec<u64> {
vec![
node.state.values.as_ptr() as u64,
node.state.table.biases().as_ptr() as u64,
node.state.table.primaries().as_ptr() as u64,
node.state.table.aliases().as_ptr() as u64,
node.state.values.len() as u64,
]
}
fn weighted_pick_decompose(node: &WeightedPick) -> DecomposedGraph {
let spec: String = node
.state
.values
.iter()
.zip(node.state.weights.iter())
.map(|(v, w)| format!("{v}:{w}"))
.collect::<Vec<_>>()
.join(";");
let mut g = DecomposedGraph::new(1);
let wu = g.add_node(
Box::new(WeightedU64::new(spec)),
vec![DecomposedWire::Input(0)],
);
g.set_outputs(vec![DecomposedWire::Node(wu, 0)]);
g
}
#[polydat::polydat_node(
category = Weighted,
compiled_u64 = weighted_pick_jit,
jit_constants = weighted_pick_jit_constants,
decompose = weighted_pick_decompose,
)]
fn weighted_pick(
input: u64,
spec: polydat::derive_support::Const<&str>,
#[poly_const(parse_weighted_pick_spec, from = spec)] state: &WeightedPickState,
) -> u64 {
let _ = spec; let idx = state.table.sample(input) as usize;
state.values[idx]
}
#[derive(Default)]
pub struct DynamicWeightedMemo {
spec: String,
values: Vec<String>,
table: Option<AliasTableU64>,
parsed: bool,
}
impl DynamicWeightedMemo {
fn select(&mut self, spec: &str, selector: u64) -> &str {
if !self.parsed || self.spec != spec {
let (values, weights) = parse_weighted_str_spec(spec);
self.table = if values.is_empty() {
None
} else {
Some(AliasTableU64::from_weights(&weights))
};
self.values = values;
self.spec.clear();
self.spec.push_str(spec);
self.parsed = true;
}
match &self.table {
Some(table) => &self.values[table.sample(selector) as usize],
None => "",
}
}
#[cfg(test)]
fn spec(&self) -> Option<&str> {
self.parsed.then_some(self.spec.as_str())
}
}
pub(crate) mod dynamic_weighted_state {
use super::{DynamicWeightedMemo, DynamicWeightedSelect};
use polydat::ast::{ScratchBuf, ScratchElem, Value};
pub(crate) fn layout(_node: &DynamicWeightedSelect) -> Vec<ScratchElem> {
vec![ScratchElem::State]
}
pub(crate) fn eval(
_node: &DynamicWeightedSelect,
scratch: &mut [ScratchBuf],
inputs: &[Value],
outputs: &mut [Value],
) {
let spec = inputs[1].to_display_string();
let memo = scratch[0]
.node_state()
.get_or_insert_with(DynamicWeightedMemo::default);
outputs[0] = Value::Str(memo.select(&spec, inputs[0].as_u64()).into());
}
}
fn dynamic_weighted_compiled(
_node: &DynamicWeightedSelect,
wire_types: &[polydat::ast::PortType],
) -> polydat::ast::CompiledSlotKit {
use polydat::ast::{PortType, ScratchBuf, ScratchElem};
let spec_ty = wire_types.get(1).copied().unwrap_or(PortType::Str);
polydat::ast::CompiledSlotKit {
scratch: vec![ScratchElem::State, ScratchElem::Str],
op: Box::new(
move |inputs: &[u64], outputs: &mut [u64], scratch: &mut [ScratchBuf]| {
let selector = inputs[0];
let spec_value;
let spec: &str =
match unsafe { polydat::compile::marshal::arg_ref(spec_ty, &inputs[1..]) } {
polydat::ast::ValueRef::Str(s) => s,
other => {
spec_value = other.to_display_string();
&spec_value
}
};
let (state, out) = scratch.split_at_mut(1);
let memo = state[0]
.node_state()
.get_or_insert_with(DynamicWeightedMemo::default);
out[0].set_str(memo.select(spec, selector));
let (ptr, len) = out[0].ptr_len();
outputs[0] = ptr;
outputs[1] = len;
},
),
}
}
#[polydat::polydat_node(
category = Weighted,
compiled_slot = dynamic_weighted_compiled,
state = dynamic_weighted_state
)]
fn dynamic_weighted_select(selector: u64, weights_spec: Config<std::sync::Arc<str>>) -> String {
DynamicWeightedMemo::default()
.select(weights_spec.0.as_ref(), selector)
.to_string()
}
#[cfg(test)]
mod tests {
use super::*;
use polydat::ast::{ConstValue, PolydatNode, Slot, Value};
use polydat::compile::fusion::FusedNode;
use xxhash_rust::xxh3::xxh3_64;
#[test]
fn weighted_strings_valid_outputs() {
let node = WeightedStrings::new("alpha:0.3;beta:0.5;gamma:0.2".to_string());
let valid = ["alpha", "beta", "gamma"];
let mut out = [Value::None];
for i in 0..1000u64 {
node.eval(&[Value::U64(xxh3_64(&i.to_le_bytes()))], &mut out);
assert!(valid.contains(&out[0].as_str()));
}
}
#[test]
fn weighted_strings_respects_weights() {
let node = WeightedStrings::new("rare:0.01;common:0.99".to_string());
let mut common_count = 0u64;
let mut out = [Value::None];
let n = 10_000u64;
for i in 0..n {
node.eval(&[Value::U64(xxh3_64(&i.to_le_bytes()))], &mut out);
if out[0].as_str() == "common" {
common_count += 1;
}
}
let ratio = common_count as f64 / n as f64;
assert!(ratio > 0.90, "common should dominate, got {ratio}");
}
#[test]
fn weighted_u64_valid_outputs() {
let node = WeightedU64::new("10:0.5;20:0.3;30:0.2".to_string());
let valid = [10u64, 20, 30];
let mut out = [Value::None];
for i in 0..1000u64 {
node.eval(&[Value::U64(xxh3_64(&i.to_le_bytes()))], &mut out);
assert!(valid.contains(&out[0].as_u64()));
}
}
#[test]
fn weighted_pick_valid_outputs() {
let node = WeightedPick::new("10:0.5;20:0.3;30:0.2".to_string());
let valid = [10u64, 20, 30];
let mut out = [Value::None];
for i in 0..1000u64 {
node.eval(&[Value::U64(xxh3_64(&i.to_le_bytes()))], &mut out);
assert!(
valid.contains(&out[0].as_u64()),
"unexpected output {} at seed {i}",
out[0].as_u64()
);
}
}
#[test]
fn weighted_pick_respects_weights() {
let node = WeightedPick::new("1:0.99;2:0.01".to_string());
let mut count_1 = 0u64;
let mut out = [Value::None];
let n = 10_000u64;
for i in 0..n {
node.eval(&[Value::U64(xxh3_64(&i.to_le_bytes()))], &mut out);
if out[0].as_u64() == 1 {
count_1 += 1;
}
}
let ratio = count_1 as f64 / n as f64;
assert!(
ratio > 0.90,
"value 1 (weight 0.99) should dominate, got {ratio}"
);
}
#[test]
fn weighted_pick_single_pair() {
let node = WeightedPick::new("42:1.0".to_string());
let mut out = [Value::None];
for i in 0..100u64 {
node.eval(&[Value::U64(i)], &mut out);
assert_eq!(out[0].as_u64(), 42);
}
}
#[test]
fn weighted_pick_equal_weights() {
let node = WeightedPick::new("10:1.0;20:1.0;30:1.0".to_string());
let mut counts = [0u64; 3];
let mut out = [Value::None];
let n = 30_000u64;
for i in 0..n {
node.eval(&[Value::U64(xxh3_64(&i.to_le_bytes()))], &mut out);
match out[0].as_u64() {
10 => counts[0] += 1,
20 => counts[1] += 1,
30 => counts[2] += 1,
v => panic!("unexpected value {v}"),
}
}
for (i, c) in counts.iter().enumerate() {
let ratio = *c as f64 / n as f64;
assert!(
ratio > 0.25 && ratio < 0.42,
"value at index {i} has ratio {ratio}, expected ~0.33"
);
}
}
#[test]
fn weighted_pick_compiled_matches_eval() {
let node = WeightedPick::new("10:0.5;20:0.3;30:0.2".to_string());
let compiled = node.compiled_u64().expect("should compile");
for i in 0..10_000u64 {
let input = xxh3_64(&i.to_le_bytes());
let mut eval_out = [Value::None];
node.eval(&[Value::U64(input)], &mut eval_out);
let mut compiled_out = [0u64];
compiled(&[input], &mut compiled_out);
assert_eq!(
eval_out[0].as_u64(),
compiled_out[0],
"eval vs compiled mismatch at seed {i}"
);
}
}
#[test]
fn weighted_pick_jit_constants_shape() {
let node = WeightedPick::new("10:0.5;20:0.3;30:0.2".to_string());
let raw = node.jit_constants();
assert_eq!(raw.len(), 5); assert_eq!(raw[4], 3);
assert_eq!(raw[0], node.state.values.as_ptr() as u64);
assert_eq!(raw[1], node.state.table.biases().as_ptr() as u64);
assert_eq!(raw[2], node.state.table.primaries().as_ptr() as u64);
assert_eq!(raw[3], node.state.table.aliases().as_ptr() as u64);
}
#[test]
fn weighted_pick_equivalence_with_weighted_u64() {
let fused = WeightedPick::new("10:0.5;20:0.3;30:0.2".to_string());
let decomposed = fused.decomposed();
for i in 0..10_000u64 {
let input = xxh3_64(&i.to_le_bytes());
let mut fused_out = [Value::None];
fused.eval(&[Value::U64(input)], &mut fused_out);
let decomposed_out = decomposed.eval(&[Value::U64(input)]);
assert_eq!(
fused_out[0].as_u64(),
decomposed_out[0].as_u64(),
"equivalence failed at seed {i}"
);
}
}
#[test]
#[should_panic(expected = "weighted_pick requires at least one entry in spec")]
fn weighted_pick_rejects_empty_spec() {
let _ = WeightedPick::new("".to_string());
}
#[test]
#[should_panic(expected = "weighted_pick: malformed entry")]
fn weighted_pick_rejects_bad_format() {
let _ = WeightedPick::new("noweight".to_string());
}
#[test]
#[should_panic(expected = "weighted_pick: weight must be a positive finite f64")]
fn weighted_pick_rejects_nonpositive_weight() {
let _ = WeightedPick::new("10:0.0;20:1.0".to_string());
}
#[test]
fn dynamic_weighted_select_basic() {
let node = DynamicWeightedSelect::new();
let spec = "alpha:0.3;beta:0.5;gamma:0.2";
let valid = ["alpha", "beta", "gamma"];
let mut out = [Value::None];
for i in 0..100u64 {
node.eval(
&[
Value::U64(xxh3_64(&i.to_le_bytes())),
Value::Str(spec.into()),
],
&mut out,
);
assert!(
valid.contains(&out[0].as_str()),
"unexpected: {}",
out[0].as_str()
);
}
}
#[test]
fn dynamic_weighted_select_follows_its_spec() {
let node = DynamicWeightedSelect::new();
let spec = "a:0.5;b:0.5";
let mut out = [Value::None];
node.eval(&[Value::U64(42), Value::Str(spec.into())], &mut out);
let first = out[0].as_str().to_string();
node.eval(&[Value::U64(42), Value::Str(spec.into())], &mut out);
assert_eq!(out[0].as_str(), first);
node.eval(&[Value::U64(42), Value::Str("x:1.0".into())], &mut out);
assert_eq!(out[0].as_str(), "x");
}
#[test]
fn dynamic_weighted_select_memoizes_in_the_state() {
use polydat::ast::ScratchBuf;
let node = DynamicWeightedSelect::new();
let mut scratch: Vec<ScratchBuf> = node
.scratch_layout()
.iter()
.map(|e| ScratchBuf::new(*e))
.collect();
assert_eq!(scratch.len(), 1);
let mut out = [Value::None];
node.eval_in(
&mut scratch,
&[Value::U64(42), Value::Str("a:0.5;b:0.5".into())],
&mut out,
);
let first = out[0].as_str().to_string();
let memo = scratch[0]
.node_state()
.get::<DynamicWeightedMemo>()
.expect("filled on the first evaluation");
assert_eq!(memo.spec(), Some("a:0.5;b:0.5"));
let table_before = memo.table.as_ref().map(|t| t as *const AliasTableU64);
node.eval_in(
&mut scratch,
&[Value::U64(42), Value::Str("a:0.5;b:0.5".into())],
&mut out,
);
assert_eq!(out[0].as_str(), first);
let memo = scratch[0]
.node_state()
.get::<DynamicWeightedMemo>()
.unwrap();
assert_eq!(
memo.table.as_ref().map(|t| t as *const AliasTableU64),
table_before,
"the same spec keeps the table it built"
);
node.eval_in(
&mut scratch,
&[Value::U64(42), Value::Str("x:1.0".into())],
&mut out,
);
assert_eq!(out[0].as_str(), "x");
let memo = scratch[0]
.node_state()
.get::<DynamicWeightedMemo>()
.unwrap();
assert_eq!(memo.spec(), Some("x:1.0"));
let cloned = scratch[0].clone();
let mut cloned = cloned;
assert!(
cloned.node_state().get::<DynamicWeightedMemo>().is_none(),
"a clone of the state starts with an empty memo"
);
}
#[test]
fn dynamic_weighted_select_compiled_form_memoizes_and_agrees() {
use polydat::ast::{PortType, ScratchBuf};
let node = DynamicWeightedSelect::new();
let kit = dynamic_weighted_compiled(&node, &[PortType::U64, PortType::Str]);
let mut scratch: Vec<ScratchBuf> =
kit.scratch.iter().map(|e| ScratchBuf::new(*e)).collect();
let mut outputs = [0u64; 2];
for (spec, selector) in [("a:0.5;b:0.5", 42u64), ("a:0.5;b:0.5", 7), ("z:1.0", 3)] {
let inputs = [selector, spec.as_ptr() as usize as u64, spec.len() as u64];
(kit.op)(&inputs, &mut outputs, &mut scratch);
let got = scratch[1].to_value();
let mut want = [Value::None];
node.eval(&[Value::U64(selector), Value::Str(spec.into())], &mut want);
assert_eq!(
got.as_str(),
want[0].as_str(),
"spec {spec}, selector {selector}"
);
assert_eq!((outputs[0], outputs[1]), scratch[1].ptr_len());
assert_eq!(
scratch[0]
.node_state()
.get::<DynamicWeightedMemo>()
.and_then(|m| m.spec()),
Some(spec)
);
}
}
#[test]
fn dynamic_weighted_select_config_wire_annotation() {
let node = DynamicWeightedSelect::new();
let meta = node.meta();
let wire_inputs = meta.wire_inputs();
assert_eq!(wire_inputs.len(), 2);
assert_eq!(wire_inputs[0].wire_cost, polydat::ast::WireCost::Data);
assert_eq!(wire_inputs[1].wire_cost, polydat::ast::WireCost::Config);
}
#[test]
fn dynamic_weighted_select_e2e_init_config() {
use polydat::dsl::events::CompileEventLog;
let source = r#"
input cycle: u64
const spec := "alpha:0.3;beta:0.7"
result := dynamic_weighted_select(hash(cycle), spec)
"#;
let mut log = CompileEventLog::new();
let _k = polydat::dsl::compile::compile_polydat_with_log(source, &mut log).unwrap();
let warnings: Vec<_> = log
.events()
.iter()
.filter(|e| {
matches!(
e,
polydat::dsl::events::CompileEvent::ConfigWireCycleWarning { .. }
)
})
.collect();
assert!(warnings.is_empty(), "init-time config should not warn");
}
#[test]
fn dynamic_weighted_select_e2e_cycle_config_warns() {
use polydat::dsl::events::CompileEventLog;
let source = r#"
input cycle: u64
spec := format_u64(hash(cycle), 10)
result := dynamic_weighted_select(hash(cycle), spec)
"#;
let mut log = CompileEventLog::new();
let _k = polydat::dsl::compile::compile_polydat_with_log(source, &mut log).unwrap();
let warnings: Vec<_> = log
.events()
.iter()
.filter(|e| {
matches!(
e,
polydat::dsl::events::CompileEvent::ConfigWireCycleWarning { .. }
)
})
.collect();
assert_eq!(
warnings.len(),
1,
"cycle-time config should warn: {warnings:?}"
);
}
#[test]
fn dynamic_weighted_select_strict_rejects_cycle_config() {
use crate::hash::Hash;
use polydat::compile::assembly::{PolydatAssembler, WireRef};
use polydat::dsl::events::CompileEventLog;
use polydat::library::convert::U64ToString;
let mut asm = PolydatAssembler::new(vec!["cycle".into()]);
asm.add_node(
"hashed",
Box::new(Hash::new()),
vec![WireRef::input("cycle")],
);
asm.add_node(
"spec",
Box::new(U64ToString::default()),
vec![WireRef::node("hashed")],
);
asm.add_node(
"dws",
Box::new(DynamicWeightedSelect::new()),
vec![
WireRef::node("hashed"), WireRef::node("spec"), ],
);
asm.add_output("result", WireRef::node("dws"));
let mut log = CompileEventLog::new();
let _kernel = asm.compile_with_log(Some(&mut log)).unwrap();
let warnings: Vec<_> = log
.events()
.iter()
.filter(|e| {
matches!(
e,
polydat::dsl::events::CompileEvent::ConfigWireCycleWarning { .. }
)
})
.collect();
assert_eq!(warnings.len(), 1, "should warn in non-strict");
let mut asm2 = PolydatAssembler::new(vec!["cycle".into()]);
asm2.add_node(
"hashed",
Box::new(Hash::new()),
vec![WireRef::input("cycle")],
);
asm2.add_node(
"spec",
Box::new(U64ToString::default()),
vec![WireRef::node("hashed")],
);
asm2.add_node(
"dws",
Box::new(DynamicWeightedSelect::new()),
vec![WireRef::node("hashed"), WireRef::node("spec")],
);
asm2.add_output("result", WireRef::node("dws"));
asm2.set_strict(true);
let result = asm2.compile();
assert!(
result.is_err(),
"strict mode should reject cycle-time config wire"
);
let msg = format!("{}", result.unwrap_err());
assert!(
msg.contains("strict") || msg.contains("config"),
"error should mention strict or config: {msg}"
);
}
#[test]
fn weighted_pick_metadata_complete() {
let node = WeightedPick::new("10:0.5;20:0.3".to_string());
let meta = node.meta();
assert_eq!(meta.name, "weighted_pick");
assert_eq!(meta.ins.len(), 2);
assert!(matches!(meta.ins[0], Slot::Wire(_)));
assert!(matches!(
&meta.ins[1],
Slot::Const {
value: ConstValue::Str(_),
..
}
));
assert_eq!(meta.outs.len(), 1);
assert_eq!(meta.wire_inputs().len(), 1);
let consts = meta.const_slots();
assert_eq!(consts.len(), 1); }
}