use std::collections::{BTreeMap, BTreeSet};
use crate::flow::{At, Node, NodeId, Values};
use crate::{Interface, Program, ResourceId, ResourceKind};
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct WordsRead {
words: BTreeMap<ResourceId, BTreeSet<u32>>,
indexed: BTreeSet<ResourceId>,
bound: BTreeMap<(u32, u32), ResourceId>,
}
impl WordsRead {
pub fn of_values(values: &Values, interface: &Interface) -> WordsRead {
WordsRead::of_nodes(values, interface, &values.reached_from(values.roots()))
}
pub fn of_nodes(values: &Values, interface: &Interface, nodes: &BTreeSet<NodeId>) -> WordsRead {
let buffers: BTreeSet<ResourceId> = interface
.resources
.iter()
.filter(|r| matches!(r.kind, ResourceKind::ConstantBuffer { .. }))
.map(|r| r.id)
.collect();
let mut out = WordsRead {
bound: interface
.resources
.iter()
.filter(|r| buffers.contains(&r.id))
.filter_map(|r| r.provenance.d3d.map(|d| ((d.space, d.register), r.id)))
.collect(),
..WordsRead::default()
};
for &id in nodes {
if let Node::Global {
resource: Some(r),
at,
} = values.node(id)
&& buffers.contains(r)
{
match at {
At::Component(word) => {
out.words.entry(*r).or_default().insert(*word);
}
At::Indexed { .. } => {
out.indexed.insert(*r);
}
}
}
}
out
}
pub fn buffers(&self) -> impl Iterator<Item = ResourceId> + '_ {
self.words.keys().copied()
}
pub fn words(&self, buffer: ResourceId) -> impl Iterator<Item = u32> + '_ {
self.words.get(&buffer).into_iter().flatten().copied()
}
pub fn rows(&self, buffer: ResourceId) -> Vec<u32> {
let mut rows: Vec<u32> = self.words(buffer).map(|w| w / 4).collect();
rows.dedup();
rows
}
pub fn lanes(&self, buffer: ResourceId, row: u32) -> Vec<u32> {
self.words
.get(&buffer)
.into_iter()
.flat_map(|words| words.range(row.saturating_mul(4)..row.saturating_mul(4).saturating_add(4)))
.map(|w| w % 4)
.collect()
}
pub fn indexed(&self, buffer: ResourceId) -> bool {
self.indexed.contains(&buffer)
}
pub fn buffer_at(&self, space: u32, register: u32) -> Option<ResourceId> {
self.bound.get(&(space, register)).copied()
}
pub fn reads_at(&self, space: u32, register: u32, offset: u32, bytes: u32) -> bool {
self.buffer_at(space, register)
.is_some_and(|buffer| self.reads(buffer, offset, bytes))
}
pub fn reads(&self, buffer: ResourceId, offset: u32, bytes: u32) -> bool {
if self.indexed(buffer) {
return true;
}
let (first, last) = (
offset / 4,
offset.saturating_add(bytes.max(1).saturating_sub(1)) / 4,
);
self.words
.get(&buffer)
.is_some_and(|words| words.range(first..=last).next().is_some())
}
}
#[derive(Clone, Debug)]
pub struct ProgramGraph {
pub values: Values,
pub interface: Interface,
pub reads: WordsRead,
}
impl Program {
pub fn graph(&self) -> Result<ProgramGraph, String> {
let values = Values::of_program(self)?;
let reads = WordsRead::of_values(&values, &self.interface);
Ok(ProgramGraph {
values,
interface: self.interface.clone(),
reads,
})
}
pub fn constant_reads(&self) -> Result<WordsRead, String> {
Values::of_program(self).map(|values| WordsRead::of_values(&values, &self.interface))
}
}