feat(ir): implement resolve functionality (untested)

crates/ir/src/lib.rs

@@ -1,13 +1,17 @@
-use std::ops::RangeInclusive;
+use std::{
+    collections::{BTreeMap, BTreeSet},
+    ops::RangeInclusive,
+};
 
+use instruction::SocketCount;
 use serde::{Deserialize, Serialize};
 
 pub mod id;
 pub mod instruction;
 pub mod semi_human;
 
-pub type Map<K, V> = ahash::AHashMap<K, V>;
-pub type Set<V> = ahash::AHashSet<V>;
+pub type Map<K, V> = std::collections::BTreeMap<K, V>;
+pub type Set<V> = std::collections::BTreeSet<V>;
 
 /// Gives you a super well typed graph IR for a given human-readable repr.
 ///
@@ -64,10 +68,127 @@ pub struct GraphIr {
     /// How the data flows forward. **Dependencies** map to **dependents** here.
     edges: Map<id::Output, Set<id::Input>>,
     /// How the data flows backward. **Dependents** map to **dependencies** here.
-    rev_edges: Map<id::Input, Set<id::Output>>,
+    rev_edges: Map<id::Input, id::Output>,
 }
 
-#[derive(Clone, Debug, PartialEq, Eq, Hash, Deserialize, Serialize)]
-pub struct Span {
-    range: RangeInclusive<usize>,
+impl GraphIr {
+    // TODO: this function, but actually the whole module, screams for tests
+    /// Returns the instruction corresponding to the given ID.
+    /// Returns [`None`] if there is no such instruction in this graph IR.
+    ///
+    /// Theoretically this could be fixed easily at the expense of some memory
+    /// by just incrementing and storing some global counter,
+    /// however, at the moment there's no compelling reason
+    /// to actually have multiple [`GraphIr`]s at one point in time.
+    /// Open an issue if that poses a problem for you.
+    #[must_use]
+    pub fn resolve<'ir>(&'ir self, id: &id::Instruction) -> Option<Instruction<'ir>> {
+        let (id, kind) = self.instructions.get_key_value(id)?;
+
+        // just try each slot and see if it's connected
+        // very crude, but it works for a proof of concept
+        let SocketCount { inputs, outputs } = kind.socket_count();
+        let socket = |id: &id::Instruction, idx| id::Socket {
+            belongs_to: id.clone(),
+            // impossible since the length is limited to a u16 already
+            #[allow(clippy::cast_possible_truncation)]
+            idx: id::SocketIdx(idx as u16),
+        };
+
+        let mut inputs_from = vec![None; inputs.into()];
+        for (idx, slot) in inputs_from.iter_mut().enumerate() {
+            let input = id::Input(socket(id, idx));
+            *slot = self.rev_edges.get(&input);
+        }
+
+        let mut outputs_to = vec![None; outputs.into()];
+        for (idx, slot) in outputs_to.iter_mut().enumerate() {
+            let output = id::Output(socket(id, idx));
+            *slot = self.edges.get(&output);
+        }
+
+        Some(Instruction {
+            id,
+            kind,
+            inputs_from,
+            outputs_to,
+        })
+    }
+
+    /// Returns the instruction this input belongs to.
+    ///
+    /// The same caveats as for [`GraphIr::resolve`] apply.
+    #[must_use]
+    pub fn owner_of_input<'ir>(&'ir self, input: &id::Input) -> Option<Instruction<'ir>> {
+        self.resolve(&input.socket().belongs_to)
+    }
+
+    /// Returns the instruction this output belongs to.
+    ///
+    /// The same caveats as for [`GraphIr::resolve`] apply.
+    #[must_use]
+    pub fn owner_of_output<'ir>(&'ir self, output: &id::Output) -> Option<Instruction<'ir>> {
+        self.resolve(&output.socket().belongs_to)
+    }
+
+    #[must_use]
+    pub fn topological_sort(&self) -> Vec<Instruction> {
+        // count how many incoming edges each vertex has
+        // chances are the BTreeMap is overkill
+        let incoming_counts: BTreeMap<_, _> = self
+            .rev_edges
+            .iter()
+            .map(|(input, _)| (self.owner_of_input(input), 1))
+            .collect();
+
+        todo!()
+    }
+}
+
+/// A full instruction in context, with its inputs and outputs.
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
+pub struct Instruction<'ir> {
+    pub id: &'ir id::Instruction,
+    pub kind: &'ir instruction::Kind,
+
+    // can't have these two public since then a user might corrupt their length
+    inputs_from: Vec<Option<&'ir id::Output>>,
+    outputs_to: Vec<Option<&'ir BTreeSet<id::Input>>>,
+}
+
+impl<'ir> Instruction<'ir> {
+    /// Where this instruction gets its inputs from.
+    ///
+    /// [`None`] means that this input is unfilled,
+    /// and must be filled before the instruction can be ran.
+    #[must_use]
+    pub fn inputs_from(&self) -> &[Option<&'ir id::Output>] {
+        &self.inputs_from
+    }
+
+    /// To whom outputs are sent. [`None`] means that this output is unused.
+    #[must_use]
+    pub fn outputs_to(&self) -> &[Option<&'ir BTreeSet<id::Input>>] {
+        &self.outputs_to
+    }
+}
+
+/// Some part referred to in source code.
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
+pub struct Span {
+    // would love to use an actual [`std::ops::RangeInclusive`], but those don't implement
+    // `PartialOrd` and `Ord` unfortunately
+    /// At which byte this span starts, inclusively.
+    pub from: usize,
+    /// At which byte this span ends, inclusively.
+    pub to: usize,
+}
+
+impl From<RangeInclusive<usize>> for Span {
+    fn from(range: RangeInclusive<usize>) -> Self {
+        Self {
+            from: *range.start(),
+            to: *range.end(),
+        }
+    }
 }