use std::collections::HashMap;
use std::sync::Mutex;
use pagesift_stat::beta::{BetaCounts, Prior};
use pagesift_stat::rng::Rng;
use pagesift_stat::{Reader, Stat, Writer};
const TABLE_VERSION: u32 = 1;
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct RenderWhy {
pub learned: bool,
pub p: f64,
pub value: f64,
pub charge: f64,
}
pub type PlanObserver = std::sync::Arc<dyn Fn(u64, bool) + Send + Sync>;
pub struct RenderPlanner {
shards: Box<[Mutex<HashMap<u64, BetaCounts>>]>,
observer: std::sync::OnceLock<PlanObserver>,
}
impl Default for RenderPlanner {
fn default() -> Self {
let n = std::thread::available_parallelism()
.map_or(1, |n| n.get())
.next_power_of_two();
Self {
shards: (0..n).map(|_| Mutex::new(HashMap::new())).collect(),
observer: std::sync::OnceLock::new(),
}
}
}
thread_local! {
static RNG: std::cell::RefCell<Rng> = std::cell::RefCell::new(Rng::from_entropy());
}
impl RenderPlanner {
fn with<R>(&self, family: &str, f: impl FnOnce(&mut HashMap<u64, BetaCounts>, u64) -> R) -> R {
let key = family_key(family);
let shard = &self.shards[(key as usize) & (self.shards.len() - 1)];
f(&mut shard.lock().unwrap_or_else(|e| e.into_inner()), key)
}
pub fn decide(
&self,
family: &str,
heuristic: bool,
http_quality: f64,
render_secs: f64,
remaining_secs: f64,
browser_price: f64,
) -> (bool, RenderWhy) {
let counts = self.with(family, |m, k| m.get(&k).copied());
let value = 1.0 - http_quality.clamp(0.0, 1.0);
let charge = render_secs * (browser_price + 1.0 / remaining_secs.max(f64::MIN_POSITIVE));
match counts.filter(|c| c.trials() > 0.0) {
None => (
heuristic,
RenderWhy {
learned: false,
p: f64::from(u8::from(heuristic)),
value,
charge,
},
),
Some(c) => {
let p = RNG.with(|r| c.posterior(Prior::UNIFORM).sample(&mut r.borrow_mut()));
(
p * value > charge,
RenderWhy {
learned: true,
p,
value,
charge,
},
)
}
}
}
pub fn record(&self, family: &str, gain: bool) {
let key = self.with(family, |m, k| {
m.entry(k).or_default().observe(gain);
k
});
if let Some(o) = self.observer.get() {
o(key, gain);
}
}
pub fn observe_with(&self, o: PlanObserver) -> bool {
self.observer.set(o).is_ok()
}
pub fn import(&self, key: u64, counts: BetaCounts) {
let shard = &self.shards[(key as usize) & (self.shards.len() - 1)];
let mut m = shard.lock().unwrap_or_else(|e| e.into_inner());
let e = m.entry(key).or_default();
if e.trials() == 0.0 {
*e = counts;
}
}
pub fn mean(&self, family: &str) -> Option<f64> {
self.with(family, |m, k| m.get(&k).copied())
.filter(|c| c.trials() > 0.0)
.map(|c| c.posterior(Prior::UNIFORM).mean())
}
pub fn encode(&self) -> Vec<u8> {
let mut w = Writer::new();
w.u32(TABLE_VERSION);
let entries: Vec<(u64, BetaCounts)> = self
.shards
.iter()
.flat_map(|s| {
let m = s.lock().unwrap_or_else(|e| e.into_inner());
m.iter().map(|(k, v)| (*k, *v)).collect::<Vec<_>>()
})
.collect();
w.u64(entries.len() as u64);
for (k, v) in &entries {
w.u64(*k);
v.encode(&mut w);
}
w.into_inner()
}
pub fn decode_merge(&self, buf: &[u8]) {
let mut r = Reader::new(buf);
let read = |r: &mut Reader<'_>| -> Option<Vec<(u64, BetaCounts)>> {
if r.u32().ok()? != TABLE_VERSION {
return None;
}
let n = r.u64().ok()?;
(0..n)
.map(|_| Some((r.u64().ok()?, BetaCounts::decode(r).ok()?)))
.collect()
};
for (k, v) in read(&mut r).unwrap_or_default() {
let shard = &self.shards[(k as usize) & (self.shards.len() - 1)];
shard
.lock()
.unwrap_or_else(|e| e.into_inner())
.entry(k)
.or_default()
.merge(&v);
}
}
}
pub fn family_key(family: &str) -> u64 {
let h = blake3::hash(family.as_bytes());
let (head, _) = h
.as_bytes()
.split_first_chunk::<8>()
.expect("BLAKE3 digest is 32 bytes");
u64::from_le_bytes(*head)
}
pub fn render_family(url: &str, client_code: bool) -> String {
format!(
"{}#{}",
family(url),
if client_code { "js" } else { "static" }
)
}
pub fn fallback_family(url: &str, engine: usize) -> String {
format!("{}#after{engine}", family(url))
}
pub fn family(raw: &str) -> String {
let Ok(url) = url::Url::parse(raw) else {
return raw.to_string();
};
let section = url
.path()
.trim_start_matches('/')
.split('/')
.next()
.unwrap_or("");
format!("{}/{}", url.origin().ascii_serialization(), section)
}
pub fn quality(doc: &pagesift_content::Document, html_bytes: usize) -> Option<f64> {
let model = crate::quality::QualityModel::builtin()?;
let weights = crate::script::ScriptWeights::builtin();
Some(model.predict(&crate::quality::features(doc, html_bytes, weights)))
}