fips_md/codegen/
util.rs

1//! Small utility structures
2
3use std::{iter, ops::Range};
4
5use anyhow::{anyhow, Result};
6
7use crate::{parser::CompileTimeConstant, runtime::{InteractionIndex, ParticleIndex, SimulationIndex}};
8
9/// Compact struct for particle, simulation and interaction indices
10#[derive(Clone, Copy)]
11pub(crate) struct IndicesRef<'a> {
12    pub particles: &'a ParticleIndex,
13    pub simulations: &'a SimulationIndex,
14    pub interactions: &'a InteractionIndex,
15}
16
17impl<'a> IndicesRef<'a> {
18    pub(crate) fn new(
19        particles: &'a ParticleIndex,
20        simulations: &'a SimulationIndex,
21        interactions: &'a InteractionIndex,
22    ) -> Self {
23        Self { particles, simulations, interactions }
24    }
25}
26
27pub(crate) fn unwrap_usize_constant(constant: &CompileTimeConstant<usize>) -> Result<usize> {
28    match constant {
29        CompileTimeConstant::Literal(value) => Ok(*value),
30        CompileTimeConstant::Identifier(name) => Err(anyhow!("Unresolved identifier {}", name)),
31        CompileTimeConstant::Substituted(value, _) => Ok(*value)
32    }
33}
34
35pub(crate) fn unwrap_f64_constant(constant: &CompileTimeConstant<f64>) -> Result<f64> {
36    match constant {
37        CompileTimeConstant::Literal(value) => Ok(*value),
38        CompileTimeConstant::Identifier(name) => Err(anyhow!("Unresolved identifier {}", name)),
39        CompileTimeConstant::Substituted(value, _) => Ok(*value)
40    }
41}
42
43#[derive(Clone, PartialEq, Eq)]
44pub(crate) struct IndexRange {
45    pub(crate) start: usize,
46    pub(crate) end: usize
47}
48
49impl IndexRange {
50    pub fn new(start: usize, end: usize) -> Self {
51        Self {
52            start, end
53        }
54    }
55
56    /// Split range into n roughly equal subranges
57    pub fn split(&self, n: usize) -> Vec<IndexRange> {
58        let len = self.end - self.start;
59        let newlen = (len + n - 1)/n;
60        (0..(n-1))
61            .map(|i| IndexRange::new(self.start+i*newlen, self.start+(i+1)*newlen))
62            .chain(iter::once(IndexRange::new(self.start+(n-1)*newlen, self.end)))
63            .collect::<Vec<_>>()
64    }
65
66    pub fn len(&self) -> usize {
67        self.end - self.start
68    }
69
70    pub fn to_range(&self, element_size: usize) -> Range<usize> {
71        (self.start * element_size)..(self.end * element_size)
72    }
73}