use super::Sim;
use crate::model::{Action, LefStates, Mobility, Position, Time};
#[derive(Copy, Clone, Debug, PartialEq)]
pub enum SampleMode {
EveryDt { dt: Time, t0: Time },
Manual,
}
impl SampleMode {
pub fn every_dt(dt: Time) -> Self {
SampleMode::EveryDt { dt, t0: 0.0 }
}
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct LefFrames {
n_lefs: usize,
time: Vec<Time>,
lefs: LefStates,
}
#[derive(Copy, Clone, Debug)]
pub struct LefSnapshot<'a> {
pub time: Time,
pub left_pos: &'a [Position],
pub right_pos: &'a [Position],
pub left_mobility: &'a [Mobility],
pub right_mobility: &'a [Mobility],
}
impl LefFrames {
pub fn new(n_lefs: usize) -> Self {
LefFrames {
n_lefs,
..Default::default()
}
}
pub fn with_capacity(n_lefs: usize, n_snapshots: usize) -> Self {
LefFrames {
n_lefs,
time: Vec::with_capacity(n_snapshots),
lefs: LefStates::with_capacity(n_lefs * n_snapshots),
}
}
pub fn reserve(&mut self, n_snapshots: usize) {
self.time.reserve(n_snapshots);
self.lefs.reserve(self.n_lefs * n_snapshots);
}
pub fn n_lefs(&self) -> usize {
self.n_lefs
}
pub fn n_snapshots(&self) -> usize {
self.time.len()
}
pub fn is_empty(&self) -> bool {
self.time.is_empty()
}
pub fn time(&self) -> &[Time] {
&self.time
}
pub fn lefs(&self) -> &LefStates {
&self.lefs
}
pub fn snapshot(&self, s: usize) -> LefSnapshot<'_> {
let rows = s * self.n_lefs..(s + 1) * self.n_lefs;
LefSnapshot {
time: self.time[s],
left_pos: &self.lefs.left_pos[rows.clone()],
right_pos: &self.lefs.right_pos[rows.clone()],
left_mobility: &self.lefs.left_mobility[rows.clone()],
right_mobility: &self.lefs.right_mobility[rows],
}
}
pub fn push(&mut self, time: Time, lefs: &LefStates) {
assert_eq!(lefs.len(), self.n_lefs, "a snapshot has one row per LEF");
self.time.push(time);
self.lefs.extend_from(lefs);
}
pub fn extend_from(&mut self, other: &LefFrames) {
assert_eq!(self.n_lefs, other.n_lefs, "n_lefs mismatch");
self.time.extend_from_slice(&other.time);
self.lefs.extend_from(&other.lefs);
}
pub fn clear(&mut self) {
self.time.clear();
self.lefs.clear();
}
pub fn into_parts(self) -> (usize, Vec<Time>, LefStates) {
(self.n_lefs, self.time, self.lefs)
}
fn push_world(&mut self, time: Time, sim: &Sim) {
assert_eq!(
self.n_lefs,
sim.n_lefs(),
"recorder built for a different sim"
);
self.time.push(time);
sim.world().push_lef_states(&mut self.lefs);
}
}
#[derive(Clone, Debug)]
pub struct Recorder {
mode: SampleMode,
frames: LefFrames,
next_k: u64,
max_snapshots: Option<usize>,
}
impl Recorder {
pub fn new(mode: SampleMode, n_lefs: usize) -> Self {
if let SampleMode::EveryDt { dt, t0 } = mode {
assert!(
dt.is_finite() && dt > 0.0,
"dt must be positive and finite, got {dt}"
);
assert!(t0.is_finite(), "t0 must be finite, got {t0}");
}
Recorder {
mode,
frames: LefFrames::new(n_lefs),
next_k: 0,
max_snapshots: None,
}
}
pub fn with_max_snapshots(mut self, n: usize) -> Self {
self.set_max_snapshots(Some(n));
self
}
pub fn set_max_snapshots(&mut self, max: Option<usize>) {
if let Some(n) = max {
assert!(n > 0, "max_snapshots must be positive");
self.frames
.reserve(n.saturating_sub(self.frames.n_snapshots()));
}
self.max_snapshots = max;
}
pub fn mode(&self) -> SampleMode {
self.mode
}
pub fn frames(&self) -> &LefFrames {
&self.frames
}
pub fn is_full(&self) -> bool {
self.max_snapshots
.is_some_and(|max| self.frames.n_snapshots() >= max)
}
pub fn take_frames(&mut self) -> LefFrames {
let n_lefs = self.frames.n_lefs;
let empty = match self.max_snapshots {
Some(max) => LefFrames::with_capacity(n_lefs, max),
None => LefFrames::new(n_lefs),
};
std::mem::replace(&mut self.frames, empty)
}
pub fn drain_into(&mut self, out: &mut LefFrames) {
out.clear();
out.n_lefs = self.frames.n_lefs;
std::mem::swap(out, &mut self.frames);
if let Some(max) = self.max_snapshots {
self.frames.reserve(max);
}
}
pub fn capture(&mut self, sim: &Sim) -> bool {
if self.is_full() {
return false;
}
self.frames.push_world(sim.time(), sim);
true
}
fn grid_point(&self, k: u64) -> Option<Time> {
match self.mode {
SampleMode::EveryDt { dt, t0 } => Some(t0 + k as f64 * dt),
SampleMode::Manual => None,
}
}
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct EventFrames {
time: Vec<f64>,
leg_id: Vec<u32>,
pos: Vec<u32>,
mob: Vec<Mobility>,
kind: Vec<u8>,
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct EventColumns {
pub time: Vec<f64>,
pub leg_id: Vec<u32>,
pub pos: Vec<u32>,
pub mob: Vec<Mobility>,
pub kind: Vec<u8>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ColumnLengthError;
impl std::fmt::Display for ColumnLengthError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "event columns have different lengths")
}
}
impl std::error::Error for ColumnLengthError {}
impl TryFrom<EventColumns> for EventFrames {
type Error = ColumnLengthError;
fn try_from(c: EventColumns) -> Result<Self, Self::Error> {
let n = c.time.len();
if [c.leg_id.len(), c.pos.len(), c.mob.len(), c.kind.len()]
.iter()
.any(|&m| m != n)
{
return Err(ColumnLengthError);
}
Ok(EventFrames {
time: c.time,
leg_id: c.leg_id,
pos: c.pos,
mob: c.mob,
kind: c.kind,
})
}
}
impl EventFrames {
pub const KIND_NAMES: [&str; 4] = ["Step", "Respawn", "Keyframe", "Evict"];
pub const KEYFRAME: u8 = 2;
pub fn new() -> Self {
Self::default()
}
pub fn reserve(&mut self, n_rows: usize) {
self.time.reserve(n_rows);
self.leg_id.reserve(n_rows);
self.pos.reserve(n_rows);
self.mob.reserve(n_rows);
self.kind.reserve(n_rows);
}
pub fn with_capacity(n_rows: usize) -> Self {
EventFrames {
time: Vec::with_capacity(n_rows),
leg_id: Vec::with_capacity(n_rows),
pos: Vec::with_capacity(n_rows),
mob: Vec::with_capacity(n_rows),
kind: Vec::with_capacity(n_rows),
}
}
pub fn n_rows(&self) -> usize {
self.time.len()
}
pub fn is_empty(&self) -> bool {
self.time.is_empty()
}
pub fn time(&self) -> &[f64] {
&self.time
}
pub fn leg_id(&self) -> &[u32] {
&self.leg_id
}
pub fn pos(&self) -> &[u32] {
&self.pos
}
pub fn mob(&self) -> &[Mobility] {
&self.mob
}
pub fn kind(&self) -> &[u8] {
&self.kind
}
pub fn extend_from(&mut self, other: &EventFrames) {
self.time.extend_from_slice(&other.time);
self.leg_id.extend_from_slice(&other.leg_id);
self.pos.extend_from_slice(&other.pos);
self.mob.extend_from_slice(&other.mob);
self.kind.extend_from_slice(&other.kind);
}
pub fn clear(&mut self) {
self.time.clear();
self.leg_id.clear();
self.pos.clear();
self.mob.clear();
self.kind.clear();
}
pub fn into_columns(self) -> EventColumns {
EventColumns {
time: self.time,
leg_id: self.leg_id,
pos: self.pos,
mob: self.mob,
kind: self.kind,
}
}
fn push_event(&mut self, time: Time, action: Action, sim: &Sim) {
let legs = match action {
Action::Step { leg } => [Some(leg), None],
Action::Respawn { lef } => sim.world().lefs[lef].legs.map(Some),
Action::Evict { leg } => {
let lef = sim.world().legs[leg].lef;
sim.world().lefs[lef].legs.map(Some)
}
};
for leg in legs.into_iter().flatten() {
let state = &sim.world().legs[leg];
self.time.push(time);
self.leg_id
.push(sim.leg_id(leg).expect("leg belongs to this sim"));
self.pos.push(state.pos);
self.mob.push(state.mobility);
self.kind.push(action.code());
}
}
fn push_keyframe(&mut self, time: Time, sim: &Sim, n_legs: usize) {
assert_eq!(
n_legs,
sim.n_legs(),
"event recorder built for a different sim"
);
let lefs = sim.world().lef_states();
let legs = lefs
.left_pos
.iter()
.zip(&lefs.left_mobility)
.zip(lefs.right_pos.iter().zip(&lefs.right_mobility));
for (lef_id, (left, right)) in legs.enumerate() {
for (side, (&pos, &mob)) in [left, right].into_iter().enumerate() {
self.time.push(time);
self.leg_id.push((2 * lef_id + side) as u32);
self.pos.push(pos);
self.mob.push(mob);
self.kind.push(Self::KEYFRAME);
}
}
}
}
#[derive(Clone, Debug)]
pub struct EventRecorder {
frames: EventFrames,
max_rows: Option<usize>,
keyframe_legs: Option<usize>,
settled: bool,
}
impl Default for EventRecorder {
fn default() -> Self {
EventRecorder {
frames: EventFrames::new(),
max_rows: None,
keyframe_legs: None,
settled: true,
}
}
}
impl EventRecorder {
pub fn new() -> Self {
Self::default()
}
pub fn with_max_rows(mut self, n: usize) -> Self {
assert!(n >= 2, "max_rows must be at least 2");
if let Some(n_legs) = self.keyframe_legs {
assert!(
n >= n_legs + 2,
"max_rows must be at least n_legs + 2 with keyframes"
);
}
self.max_rows = Some(n);
self.frames.reserve(n.saturating_sub(self.frames.n_rows()));
self
}
pub fn with_keyframes(mut self, n_legs: usize) -> Self {
if let Some(max) = self.max_rows {
assert!(
max >= n_legs + 2,
"max_rows must be at least n_legs + 2 with keyframes"
);
}
self.keyframe_legs = Some(n_legs);
self
}
fn keyframe_pending(&self) -> bool {
self.keyframe_legs.is_some() && self.frames.is_empty()
}
pub fn frames(&self) -> &EventFrames {
&self.frames
}
pub fn take_frames(&mut self) -> EventFrames {
let empty = match self.max_rows {
Some(max) => EventFrames::with_capacity(max),
None => EventFrames::new(),
};
std::mem::replace(&mut self.frames, empty)
}
pub fn drain_into(&mut self, out: &mut EventFrames) {
out.clear();
std::mem::swap(out, &mut self.frames);
if let Some(max) = self.max_rows {
self.frames.reserve(max);
}
}
}
pub trait Observer: sealed::Observe {}
impl<T: sealed::Observe + ?Sized> Observer for T {}
pub(crate) mod sealed {
use super::*;
pub trait Observe {
fn sample(&mut self, bound: Time, inclusive: bool, sim: &Sim) -> bool;
fn has_room_for_event(&self) -> bool;
fn on_event(&mut self, time: Time, action: Action, sim: &Sim);
}
impl Observe for () {
fn sample(&mut self, _: Time, _: bool, _: &Sim) -> bool {
true
}
fn has_room_for_event(&self) -> bool {
true
}
fn on_event(&mut self, _: Time, _: Action, _: &Sim) {}
}
impl Observe for Recorder {
fn sample(&mut self, bound: Time, inclusive: bool, sim: &Sim) -> bool {
let now = sim.time();
while let Some(g) = self.grid_point(self.next_k) {
let due = if inclusive { g <= bound } else { g < bound };
if !due {
break;
}
if g >= now {
if self.is_full() {
return false;
}
self.frames.push_world(g, sim);
}
self.next_k += 1;
}
true
}
fn has_room_for_event(&self) -> bool {
true
}
fn on_event(&mut self, _: Time, _: Action, _: &Sim) {}
}
impl Observe for EventRecorder {
fn sample(&mut self, bound: Time, inclusive: bool, sim: &Sim) -> bool {
let now = sim.time();
self.settled = inclusive || bound > now;
if let Some(n_legs) = self.keyframe_legs
&& self.settled
&& self.frames.is_empty()
{
self.frames.push_keyframe(now, sim, n_legs);
}
true
}
fn has_room_for_event(&self) -> bool {
let Some(max) = self.max_rows else {
return true;
};
if self.keyframe_legs.is_some() && (self.frames.is_empty() || !self.settled) {
return true;
}
self.frames.n_rows() + 2 <= max
}
fn on_event(&mut self, time: Time, action: Action, sim: &Sim) {
if self.keyframe_pending() {
return;
}
self.frames.push_event(time, action, sim);
}
}
impl<A: Observe, B: Observe> Observe for (A, B) {
fn sample(&mut self, bound: Time, inclusive: bool, sim: &Sim) -> bool {
let a = self.0.sample(bound, inclusive, sim);
let b = self.1.sample(bound, inclusive, sim);
a && b
}
fn has_room_for_event(&self) -> bool {
self.0.has_room_for_event() && self.1.has_room_for_event()
}
fn on_event(&mut self, time: Time, action: Action, sim: &Sim) {
self.0.on_event(time, action, sim);
self.1.on_event(time, action, sim);
}
}
impl<T: Observe> Observe for Option<T> {
fn sample(&mut self, bound: Time, inclusive: bool, sim: &Sim) -> bool {
self.as_mut()
.is_none_or(|o| o.sample(bound, inclusive, sim))
}
fn has_room_for_event(&self) -> bool {
self.as_ref().is_none_or(|o| o.has_room_for_event())
}
fn on_event(&mut self, time: Time, action: Action, sim: &Sim) {
if let Some(o) = self {
o.on_event(time, action, sim);
}
}
}
impl<T: Observe + ?Sized> Observe for &mut T {
fn sample(&mut self, bound: Time, inclusive: bool, sim: &Sim) -> bool {
(**self).sample(bound, inclusive, sim)
}
fn has_room_for_event(&self) -> bool {
(**self).has_room_for_event()
}
fn on_event(&mut self, time: Time, action: Action, sim: &Sim) {
(**self).on_event(time, action, sim)
}
}
}