Allow embedding diagnostics into SPIR-T as attributes.

CHANGELOG.md

@@ -34,6 +34,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased] - ReleaseDate
 
 ### Added ⭐
+- [PR#22](https://github.com/EmbarkStudios/spirt/pull/22) added `Diag` and `Attr::Diagnostics`,
+  for embedding diagnostics (errors or warnings) in SPIR-T itself
 - [PR#18](https://github.com/EmbarkStudios/spirt/pull/18) added anchor-based alignment
   to multi-version pretty-printing output (the same definitions will be kept on
   the same lines in all columns, wherever possible, to improve readability)

src/lib.rs

@@ -169,6 +169,7 @@ pub mod passes {
 pub mod spv;
 
 use smallvec::SmallVec;
+use std::borrow::Cow;
 use std::collections::BTreeSet;
 
 // HACK(eddyb) work around the lack of `FxIndex{Map,Set}` type aliases elsewhere.
@@ -290,6 +291,45 @@ pub struct AttrSetDef {
     pub attrs: BTreeSet<Attr>,
 }
 
+impl AttrSetDef {
+    pub fn push_diag(&mut self, diag: Diag) {
+        // FIXME(eddyb) seriously consider moving to `BTreeMap` (see above).
+        // HACK(eddyb) this assumes `Attr::Diagnostics` is the last of `Attr`!
+        let mut attr = if let Some(Attr::Diagnostics(_)) = self.attrs.last() {
+            self.attrs.pop_last().unwrap()
+        } else {
+            Attr::Diagnostics(OrdAssertEq(vec![]))
+        };
+        match &mut attr {
+            Attr::Diagnostics(OrdAssertEq(diags)) => diags.push(diag),
+            _ => unreachable!(),
+        }
+        self.attrs.insert(attr);
+    }
+
+    // FIXME(eddyb) should this be hidden in favor of `AttrSet::append_diag`?
+    pub fn append_diag(&self, diag: Diag) -> Self {
+        let mut new_attrs = Self {
+            attrs: self.attrs.clone(),
+        };
+        new_attrs.push_diag(diag);
+        new_attrs
+    }
+}
+
+// FIXME(eddyb) should these methods be elsewhere?
+impl AttrSet {
+    // FIXME(eddyb) should this be hidden in favor of `push_diag`?
+    // FIXME(eddyb) should these methods always take multiple values?
+    pub fn append_diag(self, cx: &Context, diag: Diag) -> Self {
+        cx.intern(cx[self].append_diag(diag))
+    }
+
+    pub fn push_diag(&mut self, cx: &Context, diag: Diag) {
+        *self = self.append_diag(cx, diag);
+    }
+}
+
 /// Any semantic or non-semantic (debuginfo) decoration/modifier, that can be
 /// *optionally* applied to some declaration/definition.
 ///
@@ -310,6 +350,99 @@ pub enum Attr {
     /// that is effectively an optimization over using `OpDecorate`.
     // FIXME(eddyb) handle flags having further operands as parameters.
     SpvBitflagsOperand(spv::Imm),
+
+    /// Can be used anywhere to record [`Diag`]nostics produced during a pass,
+    /// while allowing the pass to continue (and its output to be pretty-printed).
+    //
+    // HACK(eddyb) this is the last variant to control printing order, but also
+    // to make `push_diag`/`append_diag` above work correctly!
+    Diagnostics(OrdAssertEq<Vec<Diag>>),
+}
+
+/// Diagnostics produced by SPIR-T passes, and recorded in [`Attr::Diagnostics`].
+#[derive(Clone, PartialEq, Eq, Hash)]
+pub struct Diag {
+    pub level: DiagLevel,
+    // FIXME(eddyb) this may want to be `SmallVec` and/or `Rc`?
+    pub message: Vec<DiagMsgPart>,
+}
+
+impl Diag {
+    pub fn new(level: DiagLevel, message: impl IntoIterator<Item = DiagMsgPart>) -> Self {
+        Self {
+            level,
+            message: message.into_iter().collect(),
+        }
+    }
+
+    // FIMXE(eddyb) make macros more ergonomic than this, for interpolation.
+    #[track_caller]
+    pub fn bug(message: impl IntoIterator<Item = DiagMsgPart>) -> Self {
+        Self::new(DiagLevel::Bug(std::panic::Location::caller()), message)
+    }
+
+    pub fn err(message: impl IntoIterator<Item = DiagMsgPart>) -> Self {
+        Self::new(DiagLevel::Error, message)
+    }
+
+    pub fn warn(message: impl IntoIterator<Item = DiagMsgPart>) -> Self {
+        Self::new(DiagLevel::Warning, message)
+    }
+}
+
+/// The "severity" level of a [`Diag`]nostic.
+///
+/// Note: `Bug` diagnostics track their emission point for easier identification.
+#[derive(Copy, Clone, PartialEq, Eq, Hash)]
+pub enum DiagLevel {
+    Bug(&'static std::panic::Location<'static>),
+    Error,
+    Warning,
+}
+
+/// One part of a [`Diag`]nostic message, allowing rich interpolation.
+///
+/// Note: [`visit::Visitor`] and [`transform::Transformer`] *do not* interact
+/// with any interpolated information, and it's instead treated as "frozen" data.
+#[derive(Clone, PartialEq, Eq, Hash)]
+pub enum DiagMsgPart {
+    Plain(Cow<'static, str>),
+
+    // FIXME(eddyb) use `dyn Trait` instead of listing out a few cases.
+    Attrs(AttrSet),
+    Type(Type),
+    Const(Const),
+}
+
+// FIXME(eddyb) move this out of `lib.rs` and/or define with a macro.
+impl From<&'static str> for DiagMsgPart {
+    fn from(s: &'static str) -> Self {
+        Self::Plain(s.into())
+    }
+}
+
+impl From<String> for DiagMsgPart {
+    fn from(s: String) -> Self {
+        Self::Plain(s.into())
+    }
+}
+
+impl From<AttrSet> for DiagMsgPart {
+    fn from(attrs: AttrSet) -> Self {
+        Self::Attrs(attrs)
+    }
+}
+
+impl From<Type> for DiagMsgPart {
+    fn from(ty: Type) -> Self {
+        Self::Type(ty)
+    }
+}
+
+impl From<Const> for DiagMsgPart {
+    fn from(ct: Const) -> Self {
+        Self::Const(ct)
+    }
 }
 
 /// Wrapper to limit `Ord` for interned index types (e.g. [`InternedStr`])

src/print/mod.rs

@@ -25,10 +25,10 @@ use crate::visit::{DynVisit, InnerVisit, Visit, Visitor};
 use crate::{
     cfg, spv, AddrSpace, Attr, AttrSet, AttrSetDef, Const, ConstCtor, ConstDef, Context,
     ControlNode, ControlNodeDef, ControlNodeKind, ControlNodeOutputDecl, ControlRegion,
-    ControlRegionDef, ControlRegionInputDecl, DataInst, DataInstDef, DataInstKind, DeclDef,
-    EntityListIter, ExportKey, Exportee, Func, FuncDecl, FuncParam, FxIndexMap, GlobalVar,
-    GlobalVarDecl, GlobalVarDefBody, Import, Module, ModuleDebugInfo, ModuleDialect, SelectionKind,
-    Type, TypeCtor, TypeCtorArg, TypeDef, Value,
+    ControlRegionDef, ControlRegionInputDecl, DataInst, DataInstDef, DataInstKind, DeclDef, Diag,
+    DiagLevel, DiagMsgPart, EntityListIter, ExportKey, Exportee, Func, FuncDecl, FuncParam,
+    FxIndexMap, GlobalVar, GlobalVarDecl, GlobalVarDefBody, Import, Module, ModuleDebugInfo,
+    ModuleDialect, OrdAssertEq, SelectionKind, Type, TypeCtor, TypeCtorArg, TypeDef, Value,
 };
 use rustc_hash::FxHashMap;
 use smallvec::SmallVec;
@@ -419,6 +419,20 @@ impl<'a> Visitor<'a> for Plan<'a> {
         self.use_node(Node::ModuleDebugInfo, debug_info);
     }
 
+    fn visit_attr(&mut self, attr: &'a Attr) {
+        // HACK(eddyb) the interpolated parts aren't visited by default
+        // (as they're "inert data").
+        if let Attr::Diagnostics(OrdAssertEq(diags)) = attr {
+            for diag in diags {
+                let Diag { level, message } = diag;
+                match level {
+                    DiagLevel::Bug(_) | DiagLevel::Error | DiagLevel::Warning => {}
+                }
+                message.inner_visit_with(self);
+            }
+        }
+    }
+
     fn visit_func_decl(&mut self, func_decl: &'a FuncDecl) {
         if let DeclDef::Present(func_def_body) = &func_decl.def {
             if let Some(cfg) = &func_def_body.unstructured_cfg {
@@ -480,6 +494,27 @@ impl Plan<'_> {
         self.print(&Printer::new(self))
             .map_pretty_fragments(|fragment| fragment.layout_with_max_line_width(MAX_LINE_WIDTH))
     }
+
+    /// Like `pretty_print`, but separately pretty-printing "root dependencies"
+    /// and the "root" itself (useful for nesting pretty-printed SPIR-T elsewhere).
+    pub fn pretty_print_deps_and_root_separately(
+        &self,
+    ) -> (
+        Versions<pretty::FragmentPostLayout>,
+        Versions<pretty::FragmentPostLayout>,
+    ) {
+        let printer = Printer::new(self);
+        (
+            self.print_with_node_filter(&printer, |node| !matches!(node, Node::Root))
+                .map_pretty_fragments(|fragment| {
+                    fragment.layout_with_max_line_width(MAX_LINE_WIDTH)
+                }),
+            self.print_with_node_filter(&printer, |node| matches!(node, Node::Root))
+                .map_pretty_fragments(|fragment| {
+                    fragment.layout_with_max_line_width(MAX_LINE_WIDTH)
+                }),
+        )
+    }
 }
 
 pub struct Printer<'a> {
@@ -558,7 +593,10 @@ impl<'a> Printer<'a> {
                                             // are printed as comments outside
                                             // the `#{...}` syntax, they can't
                                             // work unless they're printed inline.
-                                            matches!(attr, Attr::SpvDebugLine { .. })
+                                            matches!(
+                                                attr,
+                                                Attr::Diagnostics(_) | Attr::SpvDebugLine { .. }
+                                            )
                                         })
                                 }
                                 CxInterned::Type(ty) => {
@@ -801,6 +839,8 @@ impl Printer<'_> {
         pretty::Styles {
             color_opacity: Some(0.3),
             size: Some(-4),
+            // FIXME(eddyb) this looks wrong for some reason?
+            // subscript: true,
             ..pretty::Styles::color(pretty::palettes::simple::DARK_GRAY)
         }
     }
@@ -1274,6 +1314,16 @@ impl Print for Func {
 impl Print for Plan<'_> {
     type Output = Versions<pretty::Fragment>;
     fn print(&self, printer: &Printer<'_>) -> Versions<pretty::Fragment> {
+        self.print_with_node_filter(printer, |_| true)
+    }
+}
+
+impl Plan<'_> {
+    fn print_with_node_filter(
+        &self,
+        printer: &Printer<'_>,
+        filter: impl Fn(&Node) -> bool,
+    ) -> Versions<pretty::Fragment> {
         let num_versions = self.per_version_name_and_node_defs.len();
         let per_node_versions_with_repeat_count = printer
             .use_styles
@@ -1282,6 +1332,7 @@ impl Print for Plan<'_> {
                 Use::Node(node) => Some(node),
                 _ => None,
             })
+            .filter(filter)
             .map(|node| -> SmallVec<[_; 1]> {
                 if num_versions == 0 {
                     return [].into_iter().collect();
@@ -1717,6 +1768,82 @@ impl Print for Attr {
     type Output = (AttrStyle, pretty::Fragment);
     fn print(&self, printer: &Printer<'_>) -> (AttrStyle, pretty::Fragment) {
         match self {
+            Attr::Diagnostics(diags) => (
+                AttrStyle::Comment,
+                pretty::Fragment::new(
+                    diags
+                        .0
+                        .iter()
+                        .map(|diag| {
+                            let Diag { level, message } = diag;
+
+                            // FIXME(eddyb) the plan was to use 💥/⮿/⚠
+                            // for bug/error/warning, but it doesn't really
+                            // render correctly, so allcaps it is for now.
+                            let (icon, icon_color) = match level {
+                                DiagLevel::Bug(_) => ("BUG", pretty::palettes::simple::MAGENTA),
+                                DiagLevel::Error => ("ERR", pretty::palettes::simple::RED),
+                                DiagLevel::Warning => ("WARN", pretty::palettes::simple::YELLOW),
+                            };
+
+                            let grayish =
+                                |[r, g, b]: [u8; 3]| [(r / 2) + 64, (g / 2) + 64, (b / 2) + 64];
+                            let comment_style = pretty::Styles::color(grayish(icon_color));
+
+                            // FIXME(eddyb) maybe make this a link to the source code?
+                            let bug_location_prefix = match level {
+                                DiagLevel::Bug(location) => {
+                                    let location = location.to_string();
+                                    let location = match location.split_once("/src/") {
+                                        Some((_path_prefix, intra_src)) => intra_src,
+                                        None => &location,
+                                    };
+                                    comment_style.clone().apply(format!("[{location}] ")).into()
+                                }
+                                DiagLevel::Error | DiagLevel::Warning => {
+                                    pretty::Fragment::default()
+                                }
+                            };
+
+                            let mut printed_message = message
+                                .print(printer)
+                                .insert_name_before_def(pretty::Fragment::default());
+
+                            // HACK(eddyb) apply the right style to all the plain
+                            // text parts of the already-printed message.
+                            for node in &mut printed_message.nodes {
+                                if let pretty::Node::Text(text) = node {
+                                    *node = comment_style.clone().apply(mem::take(text));
+                                }
+                            }
+
+                            // HACK(eddyb) this would ideally use line comments,
+                            // but adding the line prefix properly to everything
+                            // is a bit of a pain without special `pretty` support.
+                            pretty::Node::InlineOrIndentedBlock(vec![pretty::Fragment::new([
+                                comment_style.clone().apply("/* ").into(),
+                                pretty::Styles {
+                                    thickness: Some(3),
+
+                                    // HACK(eddyb) this allows larger "icons"
+                                    // without adding gaps via `line-height`.
+                                    subscript: true,
+                                    size: Some(2),
+
+                                    ..pretty::Styles::color(icon_color)
+                                }
+                                .apply(icon)
+                                .into(),
+                                " ".into(),
+                                bug_location_prefix,
+                                printed_message,
+                                comment_style.apply(" */").into(),
+                            ])])
+                        })
+                        .intersperse(pretty::Node::ForceLineSeparation),
+                ),
+            ),
+
             Attr::SpvAnnotation(spv::Inst { opcode, imms }) => {
                 struct ImplicitTargetId;
 
@@ -1765,6 +1892,25 @@ impl Print for Attr {
     }
 }
 
+impl Print for Vec<DiagMsgPart> {
+    // HACK(eddyb) this is only `AttrsAndDef` so it can be pretty-printed as a
+    // root - but ideally `DynNodeDef` should allow `pretty::Fragment` too.
+    type Output = AttrsAndDef;
+
+    fn print(&self, printer: &Printer<'_>) -> AttrsAndDef {
+        let msg = pretty::Fragment::new(self.iter().map(|part| match part {
+            DiagMsgPart::Plain(text) => pretty::Node::Text(text.clone()).into(),
+            DiagMsgPart::Attrs(attrs) => attrs.print(printer),
+            DiagMsgPart::Type(ty) => ty.print(printer),
+            DiagMsgPart::Const(ct) => ct.print(printer),
+        }));
+        AttrsAndDef {
+            attrs: pretty::Fragment::default(),
+            def_without_name: msg,
+        }
+    }
+}
+
 impl Print for TypeDef {
     type Output = AttrsAndDef;
     fn print(&self, printer: &Printer<'_>) -> AttrsAndDef {