zirium 0.4.0

//! Deterministic generic MLIR output from semantic documents.

pub(crate) mod compare;

use std::{
    collections::{HashMap, HashSet},
    fmt,
    io::{self, Write},
    path::Path,
};

use crate::dialect::DialectRegistry;
use crate::lexer::TokenKind;
use crate::semantic::{
    AffineExprValue, AttributeValue, BlockId, Document, LargeAttributeValue, LocationValue,
    MemRefLayout, OperationId, RegionId, ShapedDimension, TypeValue, ValidationError, ValueId,
    ValueReference, decode_mlir_string, quote_mlir_string,
};
use crate::{SyntaxKind, source::TextRange};

/// Whitespace policy for generated semantic output.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum PrintLayout {
    Compact,
    #[default]
    Pretty,
}

/// Controls how much metadata is retained on enclosing selection shells.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum FragmentScope {
    #[default]
    Full,
    Minimal,
}

/// Chooses generic MLIR or registered custom syntax where available.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum DialectPrintMode {
    #[default]
    GenericOnly,
    PreferCustom,
}

/// Failure while printing a semantic document.
#[derive(Debug)]
pub enum PrintError {
    IncompleteDocument,
    InvalidDocument(ValidationError),
    Format(fmt::Error),
    Io(io::Error),
    UnsafeSelection(String),
}

/// Failure while combining retained source with generated replacements.
#[derive(Debug)]
pub enum PreserveError {
    NotHybrid,
    IncompleteDocument,
    InvalidDocument(ValidationError),
    MissingSyntaxMapping,
    UnknownCustomSyntax(TextRange),
    Format(fmt::Error),
    Io(io::Error),
}

impl fmt::Display for PreserveError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::NotHybrid => f.write_str("source-preserving output requires a hybrid document"),
            Self::IncompleteDocument => {
                f.write_str("cannot regenerate an incomplete semantic document")
            }
            Self::InvalidDocument(error) => write!(f, "cannot preserve invalid document: {error}"),
            Self::MissingSyntaxMapping => {
                f.write_str("dirty semantic syntax has no retained source range")
            }
            Self::UnknownCustomSyntax(range) => write!(
                f,
                "dirty replacement at {range} contains unknown custom syntax"
            ),
            Self::Format(error) => error.fmt(f),
            Self::Io(error) => error.fmt(f),
        }
    }
}

impl std::error::Error for PreserveError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::InvalidDocument(error) => Some(error),
            Self::Format(error) => Some(error),
            Self::Io(error) => Some(error),
            _ => None,
        }
    }
}
impl fmt::Display for PrintError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::IncompleteDocument => f.write_str("cannot print an incomplete semantic document"),
            Self::InvalidDocument(error) => write!(f, "cannot print invalid document: {error}"),
            Self::Format(e) => e.fmt(f),
            Self::Io(e) => e.fmt(f),
            Self::UnsafeSelection(message) => f.write_str(message),
        }
    }
}
impl std::error::Error for PrintError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::InvalidDocument(error) => Some(error),
            Self::Format(error) => Some(error),
            Self::Io(error) => Some(error),
            Self::IncompleteDocument | Self::UnsafeSelection(_) => None,
        }
    }
}

impl Document {
    /// Formats one semantic type as valid canonical MLIR.
    pub fn canonical_type_spelling(&self, value: &TypeValue) -> Result<String, PrintError> {
        let mut spelling = String::new();
        Printer::new(
            self,
            &mut spelling,
            PrintLayout::Compact,
            DialectPrintMode::GenericOnly,
            &DialectRegistry::EMPTY,
        )
        .type_value(value)
        .map_err(PrintError::Format)?;
        Ok(spelling)
    }

    /// Formats one semantic attribute as valid canonical MLIR.
    ///
    /// This is useful when a value has no independent source spelling, such as
    /// an element reached through an attribute alias.
    pub fn canonical_attribute_spelling(
        &self,
        value: &AttributeValue,
    ) -> Result<String, PrintError> {
        let mut spelling = String::new();
        Printer::new(
            self,
            &mut spelling,
            PrintLayout::Compact,
            DialectPrintMode::GenericOnly,
            &DialectRegistry::EMPTY,
        )
        .attribute(value)
        .map_err(PrintError::Format)?;
        Ok(spelling)
    }

    /// Prints selected operations with the operation, region, and block shells
    /// that contain them. Unselected siblings are omitted.
    ///
    /// Rejects stale or foreign operation handles before writing any output.
    pub fn write_selection<W: io::Write>(
        &self,
        sink: &mut W,
        selected: &[OperationId],
        layout: PrintLayout,
        registry: &DialectRegistry,
    ) -> Result<(), PrintError> {
        self.write_selection_with_scope(sink, selected, layout, registry, FragmentScope::Full)
    }

    pub fn write_selection_with_scope<W: io::Write>(
        &self,
        sink: &mut W,
        selected: &[OperationId],
        layout: PrintLayout,
        registry: &DialectRegistry,
        scope: FragmentScope,
    ) -> Result<(), PrintError> {
        self.write_selection_with_scope_and_mode(
            sink,
            selected,
            layout,
            registry,
            scope,
            DialectPrintMode::PreferCustom,
        )
    }

    pub(crate) fn write_selection_with_scope_and_mode<W: io::Write>(
        &self,
        sink: &mut W,
        selected: &[OperationId],
        layout: PrintLayout,
        registry: &DialectRegistry,
        scope: FragmentScope,
        mode: DialectPrintMode,
    ) -> Result<(), PrintError> {
        for &operation in selected {
            self.check_operation(operation)
                .map_err(|error| PrintError::UnsafeSelection(error.to_string()))?;
        }
        let mut selected = selected.iter().copied().collect::<HashSet<_>>();
        let mut worklist = selected.iter().copied().collect::<Vec<_>>();
        while let Some(operation) = worklist.pop() {
            for &region in self.operation_regions(operation).unwrap_or(&[]) {
                for &block in self
                    .region(region)
                    .and_then(|region| region.blocks(self))
                    .unwrap_or(&[])
                {
                    for &child in self.block_operations(block).unwrap_or(&[]) {
                        if selected.insert(child) {
                            worklist.push(child);
                        }
                    }
                }
            }
        }
        let value_names = self.selection_value_names(&selected)?;
        let mut adapter = IoAdapter { sink, error: None };
        let result = Printer::new(self, &mut adapter, layout, mode, registry)
            .with_selection(&selected, value_names, scope)
            .document();
        if let Some(error) = adapter.error {
            return Err(PrintError::Io(error));
        }
        result.map_err(PrintError::Format)
    }

    fn selection_value_names(
        &self,
        selected: &HashSet<OperationId>,
    ) -> Result<HashMap<ValueId, String>, PrintError> {
        let mut opaque_names = HashSet::new();
        for &operation in selected {
            let Some(text) = self.operation_unparsed_text(operation) else {
                continue;
            };
            let source = crate::source::Source::new(text.to_vec()).map_err(|_| {
                PrintError::UnsafeSelection(
                    "cannot inspect selected opaque operation SSA names".to_owned(),
                )
            })?;
            for token in crate::lexer::lex(&source).tokens() {
                if token.kind() == TokenKind::PercentIdentifier {
                    opaque_names.insert(
                        String::from_utf8_lossy(source.slice(token.range()).unwrap()).into_owned(),
                    );
                }
            }
        }

        let mut preserved = HashMap::<ValueId, String>::new();
        for &operation in selected {
            if self.operation_unparsed_text(operation).is_none() {
                continue;
            }
            let mut parent = self
                .operation(operation)
                .and_then(|operation| operation.parent_block());
            while let Some(block) = parent {
                for argument in 0..self.block_argument_types(block).map_or(0, <[_]>::len) {
                    let Some(name) = self.block_argument_name(block, argument) else {
                        return Err(PrintError::UnsafeSelection(
                            "cannot preserve an enclosing block argument name for selected opaque syntax"
                                .to_owned(),
                        ));
                    };
                    let spelling = format!("%{name}");
                    if opaque_names.contains(&spelling) {
                        let value = ValueId::BlockArgument {
                            block,
                            argument: argument as u32,
                        };
                        preserved.insert(value, spelling);
                    }
                }
                let owner = self
                    .block(block)
                    .and_then(|block| self.region(block.parent_region()))
                    .map(|region| region.parent_operation());
                parent = owner
                    .and_then(|owner| self.operation(owner))
                    .and_then(|owner| owner.parent_block());
            }
        }

        let mut names = HashMap::new();
        let mut next = 0;
        let mut allocate = |value, names: &mut HashMap<ValueId, String>| loop {
            let spelling = format!("%v{next}");
            next += 1;
            if !opaque_names.contains(&spelling) {
                names.insert(value, spelling);
                break;
            }
        };
        for operation in self.operations() {
            for result in 0..self.result_types(operation).map_or(0, <[_]>::len) {
                allocate(
                    ValueId::OperationResult {
                        operation,
                        result: result as u32,
                    },
                    &mut names,
                );
            }
        }
        for operation in self.operations() {
            for &region in self.operation_regions(operation).unwrap_or(&[]) {
                for &block in self
                    .region(region)
                    .and_then(|region| region.blocks(self))
                    .unwrap_or(&[])
                {
                    for argument in 0..self.block_argument_types(block).map_or(0, <[_]>::len) {
                        allocate(
                            ValueId::BlockArgument {
                                block,
                                argument: argument as u32,
                            },
                            &mut names,
                        );
                    }
                }
            }
        }
        for (value, spelling) in preserved {
            names.insert(value, spelling);
        }
        Ok(names)
    }

    /// Writes deterministic generic MLIR to a [`fmt::Write`] sink.
    ///
    /// # Errors
    ///
    /// Rejects incomplete or structurally invalid documents and forwards sink
    /// formatting errors.
    pub fn print<W: fmt::Write>(
        &self,
        sink: &mut W,
        layout: PrintLayout,
    ) -> Result<(), PrintError> {
        self.preflight_print()?;
        Printer::new(
            self,
            sink,
            layout,
            DialectPrintMode::GenericOnly,
            &DialectRegistry::EMPTY,
        )
        .document()
        .map_err(PrintError::Format)
    }
    /// Alias for [`Self::print`] that names the canonical-output contract.
    ///
    /// # Errors
    ///
    /// Returns the same errors as [`Self::print`].
    pub fn print_canonical<W: fmt::Write>(
        &self,
        sink: &mut W,
        layout: PrintLayout,
    ) -> Result<(), PrintError> {
        self.print(sink, layout)
    }
    /// Returns deterministic generic MLIR as bytes.
    ///
    /// # Errors
    ///
    /// Rejects incomplete or structurally invalid documents.
    pub fn canonical_bytes(&self, layout: PrintLayout) -> Result<Vec<u8>, PrintError> {
        let mut bytes = Vec::new();
        self.write_canonical(&mut bytes, layout)?;
        Ok(bytes)
    }
    /// Writes semantic output with optional registered custom syntax.
    ///
    /// See the [registry contract](https://github.com/zayenz/zirium/blob/main/docs/custom-formats.md#registry-contract-for-each-rust-stage).
    ///
    /// In [`DialectPrintMode::PreferCustom`], a registered operation uses its
    /// built-in assembly printer or callback when that printer succeeds, and
    /// falls back to generic syntax otherwise.
    ///
    /// # Errors
    ///
    /// Rejects incomplete or structurally invalid documents and forwards sink
    /// formatting errors.
    pub fn print_with_registry<W: fmt::Write>(
        &self,
        sink: &mut W,
        layout: PrintLayout,
        mode: DialectPrintMode,
        registry: &DialectRegistry,
    ) -> Result<(), PrintError> {
        self.preflight_print()?;
        Printer::new(self, sink, layout, mode, registry)
            .document()
            .map_err(PrintError::Format)
    }
    pub fn print_io<W: io::Write>(
        &self,
        sink: &mut W,
        layout: PrintLayout,
    ) -> Result<(), PrintError> {
        self.preflight_print()?;
        let mut adapter = IoAdapter { sink, error: None };
        let result = Printer::new(
            self,
            &mut adapter,
            layout,
            DialectPrintMode::GenericOnly,
            &DialectRegistry::EMPTY,
        )
        .document();
        if let Some(error) = adapter.error {
            return Err(PrintError::Io(error));
        }
        result.map_err(PrintError::Format)
    }
    /// Writes deterministic generic MLIR derived entirely from semantic storage.
    ///
    /// # Errors
    ///
    /// Rejects incomplete or structurally invalid documents and forwards I/O
    /// errors from the sink.
    pub fn write_canonical<W: io::Write>(
        &self,
        sink: &mut W,
        layout: PrintLayout,
    ) -> Result<(), PrintError> {
        self.print_io(sink, layout)
    }
    /// Writes deterministic generic MLIR to a newly created file.
    ///
    /// # Errors
    ///
    /// Returns validation, formatting, file creation, write, or flush errors.
    pub fn print_to_file(
        &self,
        path: impl AsRef<Path>,
        layout: PrintLayout,
    ) -> Result<(), PrintError> {
        self.preflight_print()?;
        self.print_with_registry_to_preflighted_file(
            path,
            layout,
            DialectPrintMode::GenericOnly,
            &DialectRegistry::EMPTY,
        )
    }

    /// Writes semantic output with optional custom syntax to a new file.
    ///
    /// # Errors
    ///
    /// Returns validation, formatting, file creation, write, or flush errors.
    pub fn print_with_registry_to_file(
        &self,
        path: impl AsRef<Path>,
        layout: PrintLayout,
        mode: DialectPrintMode,
        registry: &DialectRegistry,
    ) -> Result<(), PrintError> {
        self.preflight_print()?;
        self.print_with_registry_to_preflighted_file(path, layout, mode, registry)
    }

    fn print_with_registry_to_preflighted_file(
        &self,
        path: impl AsRef<Path>,
        layout: PrintLayout,
        mode: DialectPrintMode,
        registry: &DialectRegistry,
    ) -> Result<(), PrintError> {
        let file = std::fs::File::create(path).map_err(PrintError::Io)?;
        let mut writer = io::BufWriter::new(file);
        let mut adapter = IoAdapter {
            sink: &mut writer,
            error: None,
        };
        let result = Printer::new(self, &mut adapter, layout, mode, registry).document();
        if let Some(error) = adapter.error {
            return Err(PrintError::Io(error));
        }
        result.map_err(PrintError::Format)?;
        writer.flush().map_err(PrintError::Io)
    }
    fn preflight_print(&self) -> Result<(), PrintError> {
        if !self.is_semantically_complete() {
            return Err(PrintError::IncompleteDocument);
        }
        self.validate().map_err(PrintError::InvalidDocument)
    }

    /// Returns source-preserving output as bytes.
    ///
    /// Unchanged ranges are copied from retained source. Dirty operations and
    /// blocks are regenerated using generic syntax. Structural insertion and
    /// erasure edits clear that retained source and mappings, so preserving
    /// output returns [`PreserveError::NotHybrid`] after either edit.
    ///
    /// # Errors
    ///
    /// Requires a complete, valid [`RetentionProfile::Hybrid`](crate::semantic::RetentionProfile::Hybrid)
    /// document with syntax mappings for every generated replacement. Replacing
    /// a range containing unknown custom syntax also fails.
    pub fn preserving_bytes(&self, layout: PrintLayout) -> Result<Vec<u8>, PreserveError> {
        self.preserving_bytes_with_mode(
            layout,
            DialectPrintMode::GenericOnly,
            &DialectRegistry::EMPTY,
        )
    }

    /// Returns source-preserving output while preferring registered custom
    /// assembly for regenerated operations and blocks.
    pub fn preserving_bytes_with_registry(
        &self,
        layout: PrintLayout,
        registry: &DialectRegistry,
    ) -> Result<Vec<u8>, PreserveError> {
        self.preserving_bytes_with_mode(layout, DialectPrintMode::PreferCustom, registry)
    }

    fn preserving_bytes_with_mode(
        &self,
        layout: PrintLayout,
        mode: DialectPrintMode,
        registry: &DialectRegistry,
    ) -> Result<Vec<u8>, PreserveError> {
        let plan = self.preserving_plan(layout, mode, registry)?;
        let source = self.source_bytes().ok_or(PreserveError::NotHybrid)?;
        let mut output = Vec::with_capacity(source.len());
        self.write_preserving_plan(&mut output, layout, mode, registry, &plan)?;
        Ok(output)
    }

    fn preserving_plan(
        &self,
        layout: PrintLayout,
        mode: DialectPrintMode,
        registry: &DialectRegistry,
    ) -> Result<Vec<PreservingReplacement>, PreserveError> {
        self.source_bytes().ok_or(PreserveError::NotHybrid)?;
        if self.retention_profile() != crate::semantic::RetentionProfile::Hybrid {
            return Err(PreserveError::NotHybrid);
        }
        if self.dirty_operations().is_empty() && self.dirty_blocks().is_empty() {
            return Ok(Vec::new());
        }
        if !self.is_semantically_complete() {
            return Err(PreserveError::IncompleteDocument);
        }
        self.validate().map_err(PreserveError::InvalidDocument)?;

        let mut replacements = Vec::new();
        for &operation in self.dirty_operations() {
            let range = self
                .operation_syntax_range(operation)
                .ok_or(PreserveError::MissingSyntaxMapping)?;
            replacements.push(PreservingReplacement::Operation(
                self.trim_trailing_trivia(range),
                operation,
            ));
        }
        for &block in self.dirty_blocks() {
            let range = self
                .block_syntax_range(block)
                .ok_or(PreserveError::MissingSyntaxMapping)?;
            replacements.push(PreservingReplacement::Block(
                self.trim_trailing_trivia(range),
                block,
            ));
        }
        replacements.sort_by_key(|replacement| {
            (
                replacement.range().start(),
                std::cmp::Reverse(replacement.range().end()),
            )
        });
        let mut planned = Vec::new();
        for replacement in replacements {
            if planned.last().is_some_and(|outer: &PreservingReplacement| {
                outer.range().end() >= replacement.range().end()
            }) {
                continue;
            }
            planned.push(replacement);
        }
        let tree = self.syntax_tree().ok_or(PreserveError::NotHybrid)?;
        let mut replacement = planned.iter().peekable();
        for custom in tree
            .subtree(tree.root())
            .into_iter()
            .flatten()
            .filter(|node| tree.kind(*node) == Some(SyntaxKind::UnparsedCustomOperation))
            .filter_map(|node| tree.text_range(node))
        {
            while replacement
                .peek()
                .is_some_and(|replacement| replacement.range().end() <= custom.start())
            {
                replacement.next();
            }
            if replacement.peek().is_some_and(|replacement| {
                let range = replacement.range();
                range.start() <= custom.start() && custom.start() < range.end()
            }) {
                return Err(PreserveError::UnknownCustomSyntax(custom));
            }
        }

        // Formatting is infallible for a validated document in practice, but run it
        // during preflight so no formatting error can occur after the first sink write.
        for replacement in &planned {
            let _ = self.render_replacement(*replacement, layout, mode, registry)?;
        }
        Ok(planned)
    }

    fn write_preserving_plan<W: io::Write>(
        &self,
        sink: &mut W,
        layout: PrintLayout,
        mode: DialectPrintMode,
        registry: &DialectRegistry,
        plan: &[PreservingReplacement],
    ) -> Result<(), PreserveError> {
        let source = self.source_bytes().ok_or(PreserveError::NotHybrid)?;
        let mut cursor = 0usize;
        for &replacement in plan {
            let range = replacement.range();
            let start = range.start() as usize;
            let end = range.end() as usize;
            sink.write_all(&source[cursor..start])
                .map_err(PreserveError::Io)?;
            let bytes = self.render_replacement(replacement, layout, mode, registry)?;
            sink.write_all(&bytes).map_err(PreserveError::Io)?;
            cursor = end;
        }
        sink.write_all(&source[cursor..]).map_err(PreserveError::Io)
    }

    /// Retains unchanged source bytes and regenerates edited regions.
    ///
    /// This mode requires [`crate::semantic::RetentionProfile::Hybrid`].
    ///
    /// # Errors
    ///
    /// Returns the planning errors described by [`Self::preserving_bytes`] or an
    /// I/O error from the sink. Planning completes before the first write.
    pub fn write_preserving<W: io::Write>(
        &self,
        sink: &mut W,
        layout: PrintLayout,
    ) -> Result<(), PreserveError> {
        let plan = self.preserving_plan(
            layout,
            DialectPrintMode::GenericOnly,
            &DialectRegistry::EMPTY,
        )?;
        self.write_preserving_plan(
            sink,
            layout,
            DialectPrintMode::GenericOnly,
            &DialectRegistry::EMPTY,
            &plan,
        )
    }

    /// Writes source-preserving output to a newly created file.
    ///
    /// # Errors
    ///
    /// Returns the planning errors described by [`Self::preserving_bytes`] or a
    /// file creation or write error.
    pub fn write_preserving_to_file(
        &self,
        path: impl AsRef<Path>,
        layout: PrintLayout,
    ) -> Result<(), PreserveError> {
        let plan = self.preserving_plan(
            layout,
            DialectPrintMode::GenericOnly,
            &DialectRegistry::EMPTY,
        )?;
        let mut file = std::fs::File::create(path).map_err(PreserveError::Io)?;
        self.write_preserving_plan(
            &mut file,
            layout,
            DialectPrintMode::GenericOnly,
            &DialectRegistry::EMPTY,
            &plan,
        )
    }

    fn render_replacement(
        &self,
        replacement: PreservingReplacement,
        layout: PrintLayout,
        mode: DialectPrintMode,
        registry: &DialectRegistry,
    ) -> Result<Vec<u8>, PreserveError> {
        match replacement {
            PreservingReplacement::Operation(_, operation) => {
                self.render_operation(operation, layout, mode, registry)
            }
            PreservingReplacement::Block(_, block) => {
                self.render_block(block, layout, mode, registry)
            }
        }
    }

    fn render_operation(
        &self,
        operation: OperationId,
        layout: PrintLayout,
        mode: DialectPrintMode,
        registry: &DialectRegistry,
    ) -> Result<Vec<u8>, PreserveError> {
        let mut output = String::new();
        Printer::new(self, &mut output, layout, mode, registry)
            .operation(operation, 0)
            .map_err(PreserveError::Format)?;
        Ok(output.into_bytes())
    }

    fn render_block(
        &self,
        block: BlockId,
        layout: PrintLayout,
        mode: DialectPrintMode,
        registry: &DialectRegistry,
    ) -> Result<Vec<u8>, PreserveError> {
        let mut output = String::new();
        Printer::new(self, &mut output, layout, mode, registry)
            .block_replacement(block)
            .map_err(PreserveError::Format)?;
        Ok(output.into_bytes())
    }

    fn block_syntax_range(&self, block: BlockId) -> Option<TextRange> {
        let operations = self.block_operations(block)?;
        let first = self.operation_syntax_range(*operations.first()?)?;
        let last = self.operation_syntax_range(*operations.last()?)?;
        let tree = self.syntax_tree()?;
        tree.subtree(tree.root())?
            .filter(|node| tree.kind(*node) == Some(SyntaxKind::Block))
            .filter_map(|node| tree.text_range(node))
            .filter(|range| range.start() <= first.start() && last.end() <= range.end())
            .min_by_key(|range| range.len())
    }

    fn trim_trailing_trivia(&self, range: TextRange) -> TextRange {
        let Some(tree) = self.syntax_tree() else {
            return range;
        };
        let tokens = tree.tokens(tree.root()).unwrap_or(&[]);
        let end_index = tokens.partition_point(|token| token.range().end() <= range.end());
        let end = tokens[..end_index]
            .iter()
            .rev()
            .take_while(|token| range.start() <= token.range().start())
            .find(|token| {
                !matches!(
                    token.kind(),
                    TokenKind::Whitespace | TokenKind::LineComment | TokenKind::Eof
                )
            })
            .map(|token| token.range().end())
            .unwrap_or(range.end());
        TextRange::new(range.start(), end).expect("trimmed source range remains ordered")
    }
}

#[derive(Clone, Copy)]
enum PreservingReplacement {
    Operation(TextRange, OperationId),
    Block(TextRange, BlockId),
}

impl PreservingReplacement {
    fn range(self) -> TextRange {
        match self {
            Self::Operation(range, _) | Self::Block(range, _) => range,
        }
    }
}

struct IoAdapter<'a, W> {
    sink: &'a mut W,
    error: Option<io::Error>,
}

impl<W: io::Write> fmt::Write for IoAdapter<'_, W> {
    fn write_str(&mut self, value: &str) -> fmt::Result {
        self.sink.write_all(value.as_bytes()).map_err(|error| {
            self.error = Some(error);
            fmt::Error
        })
    }
}

struct Printer<'a, W> {
    doc: &'a Document,
    sink: &'a mut W,
    layout: PrintLayout,
    values: HashMap<ValueId, String>,
    blocks: HashMap<BlockId, usize>,
    mode: DialectPrintMode,
    registry: &'a DialectRegistry,
    selected: Option<&'a HashSet<OperationId>>,
    fragment_scope: FragmentScope,
    custom_value_replacements: HashMap<String, String>,
}
impl<'a, W: fmt::Write> Printer<'a, W> {
    fn new(
        doc: &'a Document,
        sink: &'a mut W,
        layout: PrintLayout,
        mode: DialectPrintMode,
        registry: &'a DialectRegistry,
    ) -> Self {
        let (mut values, mut next_value) = (HashMap::new(), 0);
        for operation in doc.operations() {
            for result in 0..doc.result_types(operation).map_or(0, <[_]>::len) {
                values.insert(
                    ValueId::OperationResult {
                        operation,
                        result: result as u32,
                    },
                    format!("%v{next_value}"),
                );
                next_value += 1;
            }
        }
        let (mut blocks, mut next_block) = (HashMap::new(), 0);
        for operation in doc.operations() {
            for &region in doc.operation_regions(operation).unwrap_or(&[]) {
                for &block in doc
                    .region(region)
                    .and_then(|r| r.blocks(doc))
                    .unwrap_or(&[])
                {
                    blocks.insert(block, next_block);
                    next_block += 1;
                    for argument in 0..doc.block_argument_types(block).map_or(0, <[_]>::len) {
                        values.insert(
                            ValueId::BlockArgument {
                                block,
                                argument: argument as u32,
                            },
                            format!("%v{next_value}"),
                        );
                        next_value += 1;
                    }
                }
            }
        }
        Self {
            doc,
            sink,
            layout,
            values,
            blocks,
            mode,
            registry,
            selected: None,
            fragment_scope: FragmentScope::Full,
            custom_value_replacements: HashMap::new(),
        }
    }
    fn with_selection(
        mut self,
        selected: &'a HashSet<OperationId>,
        values: HashMap<ValueId, String>,
        fragment_scope: FragmentScope,
    ) -> Self {
        self.selected = Some(selected);
        self.fragment_scope = fragment_scope;
        // The ordinary printer already assigned canonical names. Compute the
        // selection's changes once, rather than rescanning every value for
        // every custom operation being printed.
        self.custom_value_replacements = self
            .values
            .iter()
            .filter_map(|(value, canonical)| {
                let spelling = values.get(value)?;
                (canonical != spelling).then(|| (canonical.clone(), spelling.clone()))
            })
            .collect();
        self.values = values;
        self
    }
    fn retained(&self, operation: OperationId) -> bool {
        self.selected.is_none_or(|selected| {
            selected.contains(&operation)
                || self
                    .doc
                    .operation_regions(operation)
                    .unwrap_or(&[])
                    .iter()
                    .any(|&region| {
                        self.doc
                            .region(region)
                            .and_then(|r| r.blocks(self.doc))
                            .unwrap_or(&[])
                            .iter()
                            .any(|&block| {
                                self.doc
                                    .block_operations(block)
                                    .unwrap_or(&[])
                                    .iter()
                                    .any(|&child| self.retained(child))
                            })
                    })
        })
    }
    fn rewrite_custom_value_names(&self, custom: String) -> String {
        let replacements = &self.custom_value_replacements;
        if replacements.is_empty() {
            return custom;
        }
        let source = crate::source::Source::new(custom.as_bytes().to_vec())
            .expect("a Rust string fits in Zirium source storage");
        let mut rewritten = String::with_capacity(custom.len());
        let mut cursor = 0;
        for token in crate::lexer::lex(&source).tokens() {
            if token.kind() == TokenKind::Eof {
                break;
            }
            let range = token.range().as_range();
            rewritten.push_str(&custom[cursor..range.start]);
            let spelling = &custom[range.clone()];
            if token.kind() == TokenKind::PercentIdentifier {
                rewritten.push_str(replacements.get(spelling).map_or(spelling, String::as_str));
            } else {
                rewritten.push_str(spelling);
            }
            cursor = range.end;
        }
        rewritten.push_str(&custom[cursor..]);
        rewritten
    }
    fn selected_comment(&mut self, operation: OperationId, indent: usize) -> fmt::Result {
        let Some(selected) = self.selected else {
            return Ok(());
        };
        if !selected.contains(&operation) {
            return Ok(());
        }
        let (Some(source), Some(range)) = (
            self.doc.source_bytes(),
            self.doc.operation_syntax_range(operation),
        ) else {
            return Ok(());
        };
        let start = range.start() as usize;
        let before = &source[..start];
        let line_start = before
            .iter()
            .rposition(|&b| b == b'\n')
            .map_or(0, |i| i + 1);
        if !source[line_start..start]
            .iter()
            .all(|byte| byte.is_ascii_whitespace())
        {
            return Ok(());
        }
        let mut comments = Vec::new();
        let mut cursor = line_start;
        while cursor > 0 {
            let previous_end = cursor - 1;
            let previous_start = source[..previous_end]
                .iter()
                .rposition(|&b| b == b'\n')
                .map_or(0, |i| i + 1);
            let line = &source[previous_start..previous_end];
            if String::from_utf8_lossy(line).trim_start().starts_with("//") {
                comments.push(String::from_utf8_lossy(line).trim().to_owned());
                cursor = previous_start;
            } else {
                break;
            }
        }
        comments.reverse();
        for comment in comments {
            self.sink.write_str(&comment)?;
            self.newline(indent)?;
        }
        Ok(())
    }
    fn selected_trailing_comment(&mut self, operation: OperationId) -> fmt::Result {
        let Some(selected) = self.selected else {
            return Ok(());
        };
        if !selected.contains(&operation) {
            return Ok(());
        }
        let (Some(source), Some(range)) = (
            self.doc.source_bytes(),
            self.doc
                .operation_syntax_range(operation)
                .map(|range| self.doc.trim_trailing_trivia(range)),
        ) else {
            return Ok(());
        };
        let end = range.end() as usize;
        let line_end = source[end..]
            .iter()
            .position(|&b| b == b'\n')
            .map_or(source.len(), |i| end + i);
        let tail = String::from_utf8_lossy(&source[end..line_end]);
        if let Some(comment) = tail.find("//") {
            if !tail.as_bytes()[..comment]
                .iter()
                .all(|byte| matches!(byte, b' ' | b'\t'))
            {
                return Ok(());
            }
            self.sink.write_char(' ')?;
            self.sink.write_str(tail[comment..].trim_end())?;
        }
        Ok(())
    }
    fn document(&mut self) -> fmt::Result {
        let roots = self
            .doc
            .root_operations()
            .iter()
            .copied()
            .filter(|&op| self.retained(op))
            .collect::<Vec<_>>();
        let has_roots = !roots.is_empty();
        for (index, operation) in roots.into_iter().enumerate() {
            if index != 0 {
                self.newline(0)?;
            }
            self.selected_comment(operation, 0)?;
            self.operation(operation, 0)?;
            self.selected_trailing_comment(operation)?;
        }
        if self.layout == PrintLayout::Pretty && has_roots {
            self.sink.write_char('\n')?;
        }
        Ok(())
    }
    fn operation(&mut self, id: OperationId, indent: usize) -> fmt::Result {
        if self.selected.is_some_and(|selected| selected.contains(&id))
            && let Some(source) = self.doc.operation_unparsed_text(id)
        {
            return self.sink.write_str(&String::from_utf8_lossy(source));
        }
        let results = self.doc.result_types(id).ok_or(fmt::Error)?;
        for result in 0..results.len() {
            if result != 0 {
                self.sink.write_str(", ")?;
            }
            self.value(ValueId::OperationResult {
                operation: id,
                result: result as u32,
            })?;
        }
        if !results.is_empty() {
            self.sink.write_str(" = ")?;
        }
        let minimal_shell = self.fragment_scope == FragmentScope::Minimal
            && self
                .selected
                .is_some_and(|selected| !selected.contains(&id));
        if !minimal_shell
            && self.mode == DialectPrintMode::PreferCustom
            && let Some(mut custom) = self
                .doc
                .operation_name(id)
                .and_then(|name| self.registry.operation(name))
                .and_then(|descriptor| {
                    descriptor
                        .assembly
                        .and_then(|program| program.print(self.doc, id))
                        .or_else(|| descriptor.print.and_then(|print| print(self.doc, id)))
                })
        {
            if self.selected.is_some() {
                custom = self.rewrite_custom_value_names(custom);
            }
            self.sink.write_str(&custom)?;
            let assembly = self
                .doc
                .operation_name(id)
                .and_then(|name| self.registry.operation(name))
                .and_then(|descriptor| descriptor.assembly);
            if matches!(
                assembly,
                Some(
                    crate::dialect::AssemblyProgram::BuiltinModule
                        | crate::dialect::AssemblyProgram::FuncFunc
                )
            ) {
                let regions = self.doc.operation_regions(id).ok_or(fmt::Error)?;
                if let Some(region) = regions.first() {
                    return self.region(
                        *region,
                        indent,
                        assembly == Some(crate::dialect::AssemblyProgram::FuncFunc),
                    );
                }
            }
            return Ok(());
        }
        write!(
            self.sink,
            "\"{}\"(",
            self.doc.operation_name(id).ok_or(fmt::Error)?
        )?;
        self.value_references(self.doc.operands(id).ok_or(fmt::Error)?)?;
        self.sink.write_char(')')?;
        let successors = self.doc.successors(id).ok_or(fmt::Error)?;
        if !successors.is_empty() {
            self.sink.write_str(" [")?;
            for (index, successor) in successors.iter().copied().enumerate() {
                if index != 0 {
                    self.sink.write_str(", ")?;
                }
                self.block_name(successor.block)?;
                let arguments = self.doc.successor_arguments(successor).ok_or(fmt::Error)?;
                if !arguments.is_empty() {
                    self.sink.write_str(" : (")?;
                    for (index, &argument) in arguments.iter().enumerate() {
                        if index != 0 {
                            self.sink.write_str(", ")?;
                        }
                        self.value_reference(argument)?;
                        self.sink.write_str(" : ")?;
                        let value_type = self.value_type(argument).ok_or(fmt::Error)?.clone();
                        self.type_value(&value_type)?;
                    }
                    self.sink.write_char(')')?;
                }
            }
            self.sink.write_char(']')?;
        }
        let regions = self.doc.operation_regions(id).ok_or(fmt::Error)?;
        if !regions.is_empty() {
            self.sink.write_str(" (")?;
            for (index, &region) in regions.iter().enumerate() {
                if index != 0 {
                    self.sink.write_str(", ")?;
                }
                self.region(region, indent, false)?;
            }
            self.sink.write_char(')')?;
        }
        self.dictionary(
            self.doc.operation_properties(id).ok_or(fmt::Error)?,
            " <{",
            "}>",
        )?;
        if minimal_shell {
            let structural = self
                .doc
                .operation_attributes(id)
                .ok_or(fmt::Error)?
                .iter()
                .copied()
                .filter(|(name, _)| {
                    matches!(
                        self.doc.string(*name),
                        Some(
                            "sym_name"
                                | "function_type"
                                | "sym_visibility"
                                | "arg_attrs"
                                | "res_attrs"
                                | "operand_segment_sizes"
                                | "result_segment_sizes"
                        )
                    )
                })
                .collect::<Vec<_>>();
            self.dictionary(&structural, " {", "}")?;
        } else {
            self.dictionary(
                self.doc.operation_attributes(id).ok_or(fmt::Error)?,
                " {",
                "}",
            )?;
        }
        self.sink.write_str(" : ")?;
        self.type_id(self.doc.function_type(id).ok_or(fmt::Error)?)?;
        if !minimal_shell
            && let Some(location) = self.doc.operation_location_id(id).ok_or(fmt::Error)?
        {
            self.sink.write_str(" loc(")?;
            self.location(self.doc.location_value(location).ok_or(fmt::Error)?)?;
            self.sink.write_char(')')?;
        }
        Ok(())
    }
    fn block_replacement(&mut self, block: BlockId) -> fmt::Result {
        let explicit = self.doc.block_label(block).ok_or(fmt::Error)?.is_some()
            || !self
                .doc
                .block_argument_types(block)
                .ok_or(fmt::Error)?
                .is_empty();
        if explicit {
            self.block_name(block)?;
            let arguments = self.doc.block_argument_types(block).ok_or(fmt::Error)?;
            if !arguments.is_empty() {
                self.sink.write_char('(')?;
                for (index, &ty) in arguments.iter().enumerate() {
                    if index != 0 {
                        self.sink.write_str(", ")?;
                    }
                    self.value(ValueId::BlockArgument {
                        block,
                        argument: index as u32,
                    })?;
                    self.sink.write_str(" : ")?;
                    self.type_id(ty)?;
                }
                self.sink.write_char(')')?;
            }
            self.sink.write_char(':')?;
            if !self
                .doc
                .block_operations(block)
                .ok_or(fmt::Error)?
                .is_empty()
            {
                self.newline(1)?;
            }
        }
        for (index, &operation) in self
            .doc
            .block_operations(block)
            .ok_or(fmt::Error)?
            .iter()
            .enumerate()
        {
            if index != 0 {
                self.newline(usize::from(explicit))?;
            }
            self.operation(operation, usize::from(explicit))?;
        }
        Ok(())
    }
    fn region(
        &mut self,
        id: crate::semantic::RegionId,
        indent: usize,
        implicit_entry: bool,
    ) -> fmt::Result {
        self.sink.write_char('{')?;
        let blocks = self
            .doc
            .region(id)
            .and_then(|r| r.blocks(self.doc))
            .ok_or(fmt::Error)?
            .iter()
            .copied()
            .filter(|&block| {
                self.doc
                    .block_operations(block)
                    .unwrap_or(&[])
                    .iter()
                    .any(|&op| self.retained(op))
            })
            .collect::<Vec<_>>();
        for (block_index, &block) in blocks.iter().enumerate() {
            let explicit = block_index != 0
                || (!implicit_entry
                    && (self.doc.block_label(block).ok_or(fmt::Error)?.is_some()
                        || !self
                            .doc
                            .block_argument_types(block)
                            .ok_or(fmt::Error)?
                            .is_empty()));
            if explicit {
                self.newline(indent + 1)?;
                self.block_name(block)?;
                let arguments = self.doc.block_argument_types(block).ok_or(fmt::Error)?;
                if !arguments.is_empty() {
                    self.sink.write_char('(')?;
                    for (index, &ty) in arguments.iter().enumerate() {
                        if index != 0 {
                            self.sink.write_str(", ")?;
                        }
                        self.value(ValueId::BlockArgument {
                            block,
                            argument: index as u32,
                        })?;
                        self.sink.write_str(" : ")?;
                        self.type_id(ty)?;
                    }
                    self.sink.write_char(')')?;
                }
                self.sink.write_char(':')?;
            }
            let operations = self
                .doc
                .block_operations(block)
                .ok_or(fmt::Error)?
                .iter()
                .copied()
                .filter(|&op| self.retained(op))
                .collect::<Vec<_>>();
            for operation in operations {
                self.newline(indent + 1 + usize::from(explicit))?;
                self.selected_comment(operation, indent + 1 + usize::from(explicit))?;
                self.operation(operation, indent + 1 + usize::from(explicit))?;
                self.selected_trailing_comment(operation)?;
            }
        }
        if !blocks.is_empty() {
            self.newline(indent)?;
        }
        self.sink.write_char('}')
    }
    fn dictionary(
        &mut self,
        entries: &[(u32, crate::semantic::AttributeId)],
        open: &str,
        close: &str,
    ) -> fmt::Result {
        if entries.is_empty() {
            return Ok(());
        }
        self.sink.write_str(open)?;
        for (index, &(name, value)) in entries.iter().enumerate() {
            if index != 0 {
                self.sink.write_str(", ")?;
            }
            write!(self.sink, "{} = ", self.doc.string(name).ok_or(fmt::Error)?)?;
            self.attribute(self.doc.attribute_value(value).ok_or(fmt::Error)?)?;
        }
        self.sink.write_str(close)
    }
    fn value_references(&mut self, values: &[ValueReference]) -> fmt::Result {
        for (index, &value) in values.iter().enumerate() {
            if index != 0 {
                self.sink.write_str(", ")?;
            }
            self.value_reference(value)?;
        }
        Ok(())
    }
    fn value_reference(&mut self, value: ValueReference) -> fmt::Result {
        match value {
            ValueReference::Resolved(value) => self.value(value),
            ValueReference::Invalid(_) => Err(fmt::Error),
        }
    }
    fn value(&mut self, value: ValueId) -> fmt::Result {
        self.sink
            .write_str(self.values.get(&value).ok_or(fmt::Error)?)
    }
    fn block_name(&mut self, block: BlockId) -> fmt::Result {
        write!(
            self.sink,
            "^bb{}",
            self.blocks.get(&block).ok_or(fmt::Error)?
        )
    }
    fn value_type(&self, value: ValueReference) -> Option<&TypeValue> {
        let ty = match value {
            ValueReference::Resolved(ValueId::OperationResult { operation, result }) => {
                *self.doc.result_types(operation)?.get(result as usize)?
            }
            ValueReference::Resolved(ValueId::BlockArgument { block, argument }) => *self
                .doc
                .block_argument_types(block)?
                .get(argument as usize)?,
            ValueReference::Invalid(_) => return None,
        };
        self.doc.type_value(ty)
    }
    fn type_id(&mut self, id: crate::semantic::TypeId) -> fmt::Result {
        self.type_value(self.doc.type_value(id).ok_or(fmt::Error)?)
    }
    fn type_value(&mut self, value: &TypeValue) -> fmt::Result {
        match value {
            TypeValue::Integer { width, signedness } => write!(
                self.sink,
                "{}i{width}",
                match signedness {
                    Some(true) => "s",
                    Some(false) => "u",
                    None => "",
                }
            ),
            TypeValue::Float(name) => self.sink.write_str(name),
            TypeValue::Index => self.sink.write_str("index"),
            TypeValue::Complex(element) => {
                self.sink.write_str("complex<")?;
                self.type_value(element)?;
                self.sink.write_str(" >")
            }
            TypeValue::Tuple(values) => {
                self.sink.write_str("tuple<")?;
                self.types(values)?;
                self.sink.write_str(" >")
            }
            TypeValue::Function { inputs, results } => {
                self.sink.write_char('(')?;
                self.types(inputs)?;
                self.sink.write_char(')')?;
                self.sink.write_str(" -> ")?;
                self.type_list(results)
            }
            TypeValue::Tensor {
                dimensions,
                element,
                encoding,
                unranked,
            } => {
                self.sink.write_str("tensor<")?;
                if *unranked {
                    self.sink.write_str("*x")?;
                } else {
                    self.dimensions(dimensions)?;
                }
                self.type_value(element)?;
                if let Some(encoding) = encoding {
                    self.sink.write_str(", ")?;
                    self.attribute(encoding)?;
                }
                self.sink.write_str(" >")
            }
            TypeValue::Vector {
                dimensions,
                element,
                scalable,
            } => {
                self.sink.write_str("vector<")?;
                self.dimensions_with_scalable(dimensions, scalable)?;
                self.type_value(element)?;
                self.sink.write_str(" >")
            }
            TypeValue::MemRef {
                dimensions,
                element,
                layout,
                memory_space,
            } => {
                self.sink.write_str("memref<")?;
                self.dimensions(dimensions)?;
                self.type_value(element)?;
                if let Some(layout) = layout {
                    self.sink.write_str(", ")?;
                    self.memref_layout(layout)?;
                }
                if let Some(space) = memory_space {
                    self.sink.write_str(", ")?;
                    self.attribute(space)?;
                }
                self.sink.write_str(" >")
            }
            TypeValue::Opaque(bytes) => self.opaque(bytes),
            TypeValue::Invalid(_) => Err(fmt::Error),
        }
    }
    fn type_list(&mut self, values: &[TypeValue]) -> fmt::Result {
        if values.len() == 1 {
            return self.type_value(&values[0]);
        }
        self.sink.write_char('(')?;
        self.types(values)?;
        self.sink.write_char(')')
    }
    fn types(&mut self, values: &[TypeValue]) -> fmt::Result {
        for (index, value) in values.iter().enumerate() {
            if index != 0 {
                self.sink.write_str(", ")?;
            }
            self.type_value(value)?;
        }
        Ok(())
    }
    fn dimensions(&mut self, values: &[ShapedDimension]) -> fmt::Result {
        self.dimensions_with_scalable(values, &[])
    }
    fn dimensions_with_scalable(
        &mut self,
        values: &[ShapedDimension],
        scalable: &[bool],
    ) -> fmt::Result {
        for (index, value) in values.iter().enumerate() {
            if value.scalable || scalable.get(index).copied().unwrap_or(false) {
                self.sink.write_char('[')?;
            }
            match value.size {
                Some(size) => write!(self.sink, "{size}")?,
                None => self.sink.write_char('?')?,
            }
            if value.scalable || scalable.get(index).copied().unwrap_or(false) {
                self.sink.write_char(']')?;
            }
            self.sink.write_char('x')?;
        }
        Ok(())
    }
    fn memref_layout(&mut self, value: &MemRefLayout) -> fmt::Result {
        match value {
            MemRefLayout::AffineMap(map) => self.affine_map(*map),
            MemRefLayout::Opaque {
                spelling,
                parameters,
            } => self.opaque_layout(spelling, parameters),
            MemRefLayout::Attribute(value) => self.attribute(value),
            MemRefLayout::Invalid(_) => Err(fmt::Error),
        }
    }
    fn attribute(&mut self, value: &AttributeValue) -> fmt::Result {
        match value {
            AttributeValue::Boolean(value) => {
                self.sink.write_str(if *value { "true" } else { "false" })
            }
            AttributeValue::Integer(v) | AttributeValue::Float(v) => self.sink.write_str(v),
            AttributeValue::String(v) => self.sink.write_str(
                &decode_mlir_string(v)
                    .map(|decoded| quote_mlir_string(&decoded))
                    .unwrap_or_else(|| v.clone()),
            ),
            AttributeValue::Type(v) => self.type_value(v),
            AttributeValue::Symbol(parts) => self
                .sink
                .write_str(&crate::semantic::format_symbol_path(parts, true)),
            AttributeValue::Array(values) => {
                self.sink.write_char('[')?;
                for (index, value) in values.iter().enumerate() {
                    if index != 0 {
                        self.sink.write_str(", ")?;
                    }
                    self.attribute(value)?;
                }
                self.sink.write_char(']')
            }
            AttributeValue::DenseArray {
                element_type,
                elements,
            } => {
                write!(self.sink, "array<{element_type}")?;
                if !elements.is_empty() {
                    self.sink.write_str(": ")?;
                    for (index, value) in elements.iter().enumerate() {
                        if index != 0 {
                            self.sink.write_str(", ")?;
                        }
                        self.attribute(value)?;
                    }
                }
                self.sink.write_char('>')
            }
            AttributeValue::Dictionary(values) => {
                self.sink.write_char('{')?;
                for (index, (name, value)) in values.iter().enumerate() {
                    if index != 0 {
                        self.sink.write_str(", ")?;
                    }
                    write!(self.sink, "{name} = ")?;
                    self.attribute(value)?;
                }
                self.sink.write_char('}')
            }
            AttributeValue::Location(v) => {
                self.sink.write_str("loc(")?;
                self.location(v)?;
                self.sink.write_char(')')
            }
            AttributeValue::AffineMap(map) => self.affine_map(*map),
            AttributeValue::IntegerSet(set) => self.integer_set(*set),
            AttributeValue::Large(
                LargeAttributeValue::Dense(v)
                | LargeAttributeValue::Sparse(v)
                | LargeAttributeValue::Resource(v),
            )
            | AttributeValue::WideNumber(v)
            | AttributeValue::Opaque(v) => self.opaque(v),
            AttributeValue::Invalid(_) => Err(fmt::Error),
        }
    }
    fn opaque_layout(&mut self, spelling: &str, parameters: &[AttributeValue]) -> fmt::Result {
        if parameters.is_empty() {
            return self.sink.write_str(spelling);
        }
        let bytes = spelling.as_bytes();
        let mut start = 0;
        let mut parameter = 0;
        let mut index = 0;
        while index < bytes.len() {
            if matches!(bytes[index], b'#' | b'!') {
                let token_start = index;
                index += 1;
                while index < bytes.len()
                    && (bytes[index].is_ascii_alphanumeric()
                        || matches!(bytes[index], b'_' | b'$' | b'.' | b'-'))
                {
                    index += 1;
                }
                self.sink.write_str(&spelling[start..token_start])?;
                if let Some(value) = parameters.get(parameter) {
                    self.attribute(value)?;
                    parameter += 1;
                } else {
                    self.sink.write_str(&spelling[token_start..index])?;
                }
                start = index;
            } else {
                index += 1;
            }
        }
        self.sink.write_str(&spelling[start..])
    }
    fn affine_map(&mut self, id: crate::semantic::AffineMapId) -> fmt::Result {
        let map = self.doc.affine_map(id).ok_or(fmt::Error)?;
        self.sink.write_str("affine_map<(")?;
        self.affine_vars('d', map.dimensions)?;
        self.sink.write_char(')')?;
        if map.symbols != 0 {
            self.sink.write_char('[')?;
            self.affine_vars('s', map.symbols)?;
            self.sink.write_char(']')?;
        }
        self.sink.write_str(" -> (")?;
        for (index, &expr) in map.results.iter().enumerate() {
            if index != 0 {
                self.sink.write_str(", ")?;
            }
            self.affine_expr(expr)?;
        }
        self.sink.write_str(")>")
    }
    fn integer_set(&mut self, id: crate::semantic::IntegerSetId) -> fmt::Result {
        use crate::semantic::IntegerSetRelation::*;
        let set = self.doc.integer_set(id).ok_or(fmt::Error)?;
        self.sink.write_str("affine_set<(")?;
        self.affine_vars('d', set.dimensions)?;
        self.sink.write_char(')')?;
        if set.symbols != 0 {
            self.sink.write_char('[')?;
            self.affine_vars('s', set.symbols)?;
            self.sink.write_char(']')?;
        }
        self.sink.write_str(" : (")?;
        for (index, constraint) in set.constraints.iter().enumerate() {
            if index != 0 {
                self.sink.write_str(", ")?;
            }
            self.affine_expr(constraint.left)?;
            self.sink.write_str(match constraint.relation {
                Equal => " == ",
                GreaterEqual => " >= ",
                LessEqual => " <= ",
                Invalid(_) => return Err(fmt::Error),
            })?;
            self.affine_expr(constraint.right)?;
        }
        self.sink.write_str(")>")
    }
    fn affine_vars(&mut self, prefix: char, count: u32) -> fmt::Result {
        for index in 0..count {
            if index != 0 {
                self.sink.write_str(", ")?;
            }
            write!(self.sink, "{prefix}{index}")?;
        }
        Ok(())
    }
    fn affine_expr(&mut self, id: crate::semantic::AffineExprId) -> fmt::Result {
        self.affine_expr_prec(id, 0, false)
    }
    fn affine_expr_prec(
        &mut self,
        id: crate::semantic::AffineExprId,
        parent_precedence: u8,
        right_child: bool,
    ) -> fmt::Result {
        use crate::semantic::{AffineBinaryOperator::*, AffineExprValue::*};
        match self.doc.affine_expression(id).ok_or(fmt::Error)? {
            Dimension(i) => write!(self.sink, "d{i}"),
            Symbol(i) => write!(self.sink, "s{i}"),
            Constant(v) => write!(self.sink, "{v}"),
            Binary {
                operator,
                left,
                right,
            } => {
                if matches!(operator, Subtract)
                    && matches!(self.doc.affine_expression(*left), Some(Constant(0)))
                {
                    self.sink.write_str("- ")?;
                    return self.affine_expr_prec(*right, 3, true);
                }
                let precedence = match operator {
                    Add | Subtract => 1,
                    Multiply | FloorDiv | CeilDiv | Mod => 2,
                };
                let parentheses = precedence < parent_precedence
                    || (right_child && precedence == parent_precedence);
                if parentheses {
                    self.sink.write_char('(')?;
                }
                self.affine_expr_prec(*left, precedence, false)?;
                self.sink.write_str(match operator {
                    Add => " + ",
                    Subtract => " - ",
                    Multiply => " * ",
                    FloorDiv => " floordiv ",
                    CeilDiv => " ceildiv ",
                    Mod => " mod ",
                })?;
                self.affine_expr_prec(*right, precedence, true)?;
                if parentheses {
                    self.sink.write_char(')')?;
                }
                Ok(())
            }
            Invalid(_) => Err(fmt::Error),
        }
    }
    fn location(&mut self, value: &LocationValue) -> fmt::Result {
        match value {
            LocationValue::Unknown => self.sink.write_str("unknown"),
            LocationValue::FileLineColumn { file, line, column } => {
                write!(self.sink, "{file}:{line}:{column}")
            }
            LocationValue::Name {
                name,
                child,
                metadata,
            } => {
                self.sink.write_str(name)?;
                if let Some(child) = child {
                    self.sink.write_char('(')?;
                    self.location(child)?;
                    self.sink.write_char(')')?;
                }
                if let Some(metadata) = metadata {
                    self.sink.write_str(metadata)?;
                }
                Ok(())
            }
            LocationValue::CallSite { callee, caller } => {
                self.sink.write_str("callsite(")?;
                self.location(callee)?;
                self.sink.write_str(" at ")?;
                self.location(caller)?;
                self.sink.write_char(')')
            }
            LocationValue::Fused {
                metadata,
                locations,
            } => {
                self.sink.write_str("fused")?;
                if let Some(metadata) = metadata {
                    write!(self.sink, "<{metadata}>")?;
                }
                self.sink.write_char('[')?;
                for (index, location) in locations.iter().enumerate() {
                    if index != 0 {
                        self.sink.write_str(", ")?;
                    }
                    self.location(location)?;
                }
                self.sink.write_char(']')
            }
            LocationValue::Invalid(_) => Err(fmt::Error),
        }
    }
    fn opaque(&mut self, bytes: &[u8]) -> fmt::Result {
        self.sink
            .write_str(std::str::from_utf8(bytes).map_err(|_| fmt::Error)?)
    }
    fn newline(&mut self, indent: usize) -> fmt::Result {
        match self.layout {
            PrintLayout::Compact => self.sink.write_char(' '),
            PrintLayout::Pretty => {
                self.sink.write_char('\n')?;
                for _ in 0..indent {
                    self.sink.write_str("  ")?;
                }
                Ok(())
            }
        }
    }
}