use crate::error::RenderError;
use molgfx_core::{VisualPipeline, VisualProgram, VisualStage, visual_pipeline};
use molgfx_gpu::Device;
use std::collections::HashMap;
use std::hash::{Hash, Hasher};
pub(crate) const MAX_SPECIALIZED_PIPELINES: usize = 64;
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
pub(crate) struct VisualFamily(&'static str);
impl VisualFamily {
#[must_use]
pub(crate) const fn new(name: &'static str) -> Self {
Self(name)
}
#[must_use]
pub(crate) const fn name(self) -> &'static str {
self.0
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub(crate) struct SpecializationKey {
fingerprint: u64,
stage: VisualStage,
family: VisualFamily,
}
impl SpecializationKey {
#[must_use]
pub(crate) const fn new(fingerprint: u64, stage: VisualStage, family: VisualFamily) -> Self {
Self {
fingerprint,
stage,
family,
}
}
#[must_use]
pub(crate) fn label(self) -> String {
format!(
"{:016x}:{:?}:{}",
self.fingerprint,
self.stage,
self.family.name()
)
}
}
impl Hash for SpecializationKey {
fn hash<H: Hasher>(&self, state: &mut H) {
self.fingerprint.hash(state);
(self.stage as u8).hash(state);
self.family.hash(state);
}
}
#[derive(Debug)]
pub(crate) enum Specialization<'a, D: Device> {
Ready(&'a D::Pipeline),
Interpreter {
state: VisualPipeline,
reason: &'a str,
},
}
impl<'a, D: Device> Specialization<'a, D> {
#[must_use]
pub(crate) const fn pipeline(&self) -> Option<&'a D::Pipeline> {
match self {
Self::Ready(pipeline) => Some(pipeline),
Self::Interpreter { .. } => None,
}
}
#[must_use]
pub(crate) const fn state(&self) -> VisualPipeline {
match self {
Self::Ready(_) => VisualPipeline::Specialized,
Self::Interpreter { state, .. } => *state,
}
}
#[must_use]
pub(crate) fn report(&self, key: SpecializationKey) -> String {
let reason = match self {
Self::Ready(_) => VisualPipeline::Specialized.default_reason().to_owned(),
Self::Interpreter { reason, .. } => (*reason).to_owned(),
};
self.state().report(&key.label(), &reason)
}
}
#[derive(Debug)]
struct Entry<D: Device> {
pipeline: D::Pipeline,
last_use: u64,
}
#[derive(Debug)]
pub(crate) struct SpecializedPipelines<D: Device> {
entries: HashMap<SpecializationKey, Entry<D>>,
failed: HashMap<SpecializationKey, (VisualPipeline, String)>,
capacity: usize,
clock: u64,
}
impl<D: Device> Default for SpecializedPipelines<D> {
fn default() -> Self {
Self::new()
}
}
impl<D: Device> SpecializedPipelines<D> {
#[must_use]
pub(crate) fn new() -> Self {
Self::with_capacity(MAX_SPECIALIZED_PIPELINES)
}
#[must_use]
pub(crate) fn with_capacity(capacity: usize) -> Self {
Self {
entries: HashMap::new(),
failed: HashMap::new(),
capacity: capacity.max(1),
clock: 0,
}
}
#[cfg(test)]
#[must_use]
pub(crate) fn len(&self) -> usize {
self.entries.len()
}
#[cfg(test)]
#[must_use]
pub(crate) fn failures(&self) -> usize {
self.failed.len()
}
#[cfg(test)]
#[must_use]
pub(crate) fn is_ready(&self, key: SpecializationKey) -> bool {
self.entries.contains_key(&key)
}
#[cfg(test)]
#[must_use]
pub(crate) fn failure_reason(&self, key: SpecializationKey) -> Option<&str> {
self.failed.get(&key).map(|(_, reason)| reason.as_str())
}
#[cfg(test)]
pub(crate) fn acquire(
&mut self,
key: SpecializationKey,
program: &VisualProgram,
build: impl FnOnce() -> Result<D::Pipeline, RenderError>,
) -> Specialization<'_, D> {
self.settle(key, program, build);
self.resolve(key)
}
pub(crate) fn settle(
&mut self,
key: SpecializationKey,
program: &VisualProgram,
build: impl FnOnce() -> Result<D::Pipeline, RenderError>,
) {
let (eligibility, reason) = visual_pipeline(program);
if eligibility.is_interpreter() {
let _ = self.failed.entry(key).or_insert((eligibility, reason));
return;
}
self.clock += 1;
if let Some(entry) = self.entries.get_mut(&key) {
entry.last_use = self.clock;
return;
}
if self.failed.contains_key(&key) {
return;
}
match build() {
Ok(pipeline) => {
self.make_room();
let _ = self.entries.insert(
key,
Entry {
pipeline,
last_use: self.clock,
},
);
}
Err(error) => {
let _ = self
.failed
.entry(key)
.or_insert((VisualPipeline::Failed, error.to_string()));
}
}
}
pub(crate) fn resolve(&self, key: SpecializationKey) -> Specialization<'_, D> {
if let Some(entry) = self.entries.get(&key) {
return Specialization::Ready(&entry.pipeline);
}
let (state, reason) = self.failed.get(&key).map_or(
(VisualPipeline::Failed, "not-compiled"),
|(state, reason)| (*state, reason.as_str()),
);
Specialization::Interpreter { state, reason }
}
fn make_room(&mut self) {
if self.entries.len() < self.capacity {
return;
}
let victim = self
.entries
.iter()
.min_by_key(|(_, entry)| entry.last_use)
.map(|(key, _)| *key);
if let Some(victim) = victim {
let _ = self.entries.remove(&victim);
}
}
}