uiua 0.18.1

use std::{
    collections::{BTreeMap, HashMap},
    fmt,
    hash::{Hash, Hasher},
    ops::{Index, IndexMut},
    path::PathBuf,
    str::FromStr,
    sync::Arc,
};

use dashmap::DashMap;
use ecow::{EcoString, EcoVec, eco_vec};
use indexmap::IndexMap;
use rapidhash::quality::RapidHasher;
use serde::*;

use crate::{
    BindingCounts, CodeSpan, FunctionId, Ident, InputSrc, Inputs, LocalNames, Node,
    SUBSCRIPT_DIGITS, SigNode, Signature, Sp, Span, Uiua, UiuaResult, Value,
    ast::Word,
    compile::{LocalIndex, Module},
    is_ident_char,
};

/// A compiled Uiua assembly
#[derive(Clone)]
pub struct Assembly {
    /// The top-level node
    pub root: Node,
    /// A list of dependency paths and their hashes
    pub dependencies: EcoVec<(PathBuf, u64)>,
    /// A list of top-level names
    pub exports: Arc<IndexMap<Ident, usize>>,
    /// Functions
    pub functions: EcoVec<Node>,
    /// Unexpanded index macros
    pub index_macros: Arc<IndexMap<usize, IndexMacro>>,
    /// Unexpanded code macros
    pub code_macros: Arc<IndexMap<usize, CodeMacro>>,
    /// A list of global bindings
    pub bindings: EcoVec<BindingInfo>,
    /// Indexable list of spans
    pub spans: EcoVec<Span>,
    /// Inputs used to build the assembly
    pub inputs: Inputs,
    pub(crate) dynamic_functions: EcoVec<DynFn>,
    pub(crate) test_assert_count: usize,
    /// Height of the stack on each line
    pub(crate) line_sigs: BTreeMap<u16, Signature>,
}

/// A Uiua function
///
/// This does not actually contain the function's code.
/// It is a lightweight handle that can be used to look up the function's code in an [`Assembly`].
///
/// It also contains the function's [`FunctionId`] and [`Signature`].
#[derive(Clone)]
pub struct Function {
    /// The function's id
    pub id: FunctionId,
    /// The function's signature
    pub sig: Signature,
    pub(crate) index: usize,
    hash: u64,
    pub(crate) origin: FunctionOrigin,
}

impl fmt::Debug for Function {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} ← {}", self.id, self.sig)
    }
}

impl PartialEq for Function {
    fn eq(&self, other: &Self) -> bool {
        self.id == other.id && self.sig == other.sig && self.hash == other.hash
    }
}

impl Eq for Function {}

impl Hash for Function {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.hash.hash(state);
    }
}

impl Serialize for Function {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        (&self.id, &self.sig, &self.index, &self.hash, &self.origin).serialize(serializer)
    }
}

impl<'de> Deserialize<'de> for Function {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let (id, sig, index, hash, origin) =
            <(FunctionId, Signature, usize, u64, FunctionOrigin)>::deserialize(deserializer)?;
        Ok(Function {
            id,
            sig,
            index,
            hash,
            origin,
        })
    }
}

/// Where a function came from
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum FunctionOrigin {
    /// The function was inserted by Rust code that uses the Uiua crate
    #[serde(rename = "dyn")]
    Dynamic,
    /// Function function was generated by a macro
    #[serde(rename = "macro")]
    Macro,
    /// The function is from a normal binding
    #[serde(untagged)]
    Binding(usize),
}

impl From<usize> for FunctionOrigin {
    fn from(value: usize) -> Self {
        FunctionOrigin::Binding(value)
    }
}

impl FunctionOrigin {
    /// Get the index of the function's binding
    pub fn binding(self) -> Option<usize> {
        match self {
            FunctionOrigin::Binding(i) => Some(i),
            _ => None,
        }
    }
}

/// An index macro
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IndexMacro {
    /// The words of the macro
    pub words: Vec<Sp<Word>>,
    /// Map of spans of identifiers used in the macro that were in scope
    /// when the macro was declared to their local indices. This is used
    /// for name resolution. It is keyed by span rather than by name so
    /// that names in both the declaration and invocation's scope can
    /// be disambiguated.
    pub locals: EcoVec<(CodeSpan, usize)>,
    /// The provided signature
    pub sig: Option<Signature>,
    /// Whether the macro is recursive
    pub recursive: bool,
}

/// A code macro
#[derive(Clone, Serialize, Deserialize)]
pub struct CodeMacro {
    /// The root node of the macro
    pub root: SigNode,
    /// The names that are in scope for the macro
    pub names: Arc<LocalNames>,
}

impl Assembly {
    /// Get the [`SigNode`] for a function
    pub fn sig_node(&self, f: &Function) -> SigNode {
        SigNode::new(f.sig, self[f].clone())
    }
    /// Add a function to the assembly
    pub fn add_function(
        &mut self,
        id: FunctionId,
        sig: Signature,
        mut root: Node,
        origin: impl Into<FunctionOrigin>,
    ) -> Function {
        root.optimize_early();
        let mut hasher = RapidHasher::new(1);
        root.hash(&mut hasher);
        let hash = hasher.finish();
        self.functions.push(root);
        let index = self.functions.len() - 1;
        Function {
            id,
            sig,
            index,
            hash,
            origin: origin.into(),
        }
    }
    pub(crate) fn add_binding_at(
        &mut self,
        local: LocalIndex,
        kind: BindingKind,
        span: Option<CodeSpan>,
        meta: BindingMeta,
    ) {
        let binding = BindingInfo {
            public: local.public,
            kind,
            span: span.unwrap_or_else(CodeSpan::dummy),
            meta,
            used: local.public,
        };
        if local.index < self.bindings.len() {
            self.bindings.make_mut()[local.index] = binding;
        } else {
            while self.bindings.len() < local.index {
                self.bindings.push(BindingInfo {
                    kind: BindingKind::Const(None),
                    public: false,
                    span: CodeSpan::dummy(),
                    meta: BindingMeta::default(),
                    used: true,
                });
            }
            self.bindings.push(binding);
        }
    }
    pub(crate) fn bind_const(
        &mut self,
        local: LocalIndex,
        value: Option<Value>,
        span: usize,
        meta: BindingMeta,
    ) {
        let span = self.spans[span].clone();
        self.add_binding_at(local, BindingKind::Const(value), span.code(), meta);
    }
    pub(crate) fn module(&self) -> Module {
        let mut module = Module::default();
        for (name, &index) in &*self.exports {
            let public = self.bindings[index].public;
            (module.names).insert(name.clone(), LocalIndex { index, public });
        }
        module
    }
    /// Parse a `.uasm` file into an assembly
    pub fn from_uasm(src: &str) -> Result<Self, String> {
        let rest = src;
        let (root_src, rest) = rest.split_once("DEPENDENCIES").ok_or("No dependencies")?;
        let (dependencies_src, rest) = rest.split_once("EXPORTS").ok_or("No exports")?;
        let (exports_src, rest) = rest.split_once("BINDINGS").ok_or("No bindings")?;
        let (bindings_src, rest) = rest.trim().split_once("FUNCTIONS").ok_or("No functions")?;
        let (functions_src, rest) = (rest.trim())
            .split_once("INDEX MACROS")
            .ok_or("No index macros")?;
        let (index_macros_src, rest) = (rest.trim())
            .split_once("CODE MACROS")
            .ok_or("No code macros")?;
        let (code_macros_src, rest) = rest.trim().split_once("SPANS").ok_or("No spans")?;
        let (spans_src, rest) = rest.trim().split_once("FILES").ok_or("No files")?;
        let (files_src, rest) = rest
            .trim()
            .split_once("MACRO EXPANSIONS")
            .ok_or("No macro expansions")?;
        let (expansions_src, rest) = rest
            .trim()
            .split_once("STRING INPUTS")
            .unwrap_or((rest, ""));
        let strings_src = rest.trim();

        // Root node
        let mut root = Node::empty();
        for line in root_src.lines().filter(|line| !line.trim().is_empty()) {
            let node: Node = serde_json::from_str(line).unwrap();
            root.push(node);
        }

        // Dependencies
        let mut dependencies = EcoVec::new();
        for line in dependencies_src
            .lines()
            .filter(|line| !line.trim().is_empty())
        {
            let (path, hash) = line.rsplit_once(' ').ok_or("Missing dependency hash")?;
            dependencies.push((
                PathBuf::from(path),
                hash.parse::<u64>().map_err(|_| "Invalid dependency hash")?,
            ));
        }

        // Exports
        let mut exports = IndexMap::new();
        for line in exports_src.lines().filter(|line| !line.trim().is_empty()) {
            let mut words = line.split_whitespace();
            let name = words.next().ok_or("Missing export name")?;
            let index = words
                .next()
                .ok_or("Missing export index")?
                .parse::<usize>()
                .map_err(|e| format!("Invalid export index: {e}"))?;
            exports.insert(name.into(), index);
        }
        let exports = Arc::new(exports);

        // Files
        let files = DashMap::new();
        let mut file_paths = Vec::new();
        for line in files_src.lines().filter(|line| !line.trim().is_empty()) {
            let (path, src) = line.split_once(": ").ok_or("No path")?;
            file_paths.push(path);
            let path = PathBuf::from(path);
            let src: EcoString = serde_json::from_str(src).map_err(|e| e.to_string())?;
            files.insert(path, src);
        }

        // Macro expansions
        let macros = DashMap::new();
        let mut macro_spans = Vec::new();
        for line in expansions_src
            .lines()
            .filter(|line| !line.trim().is_empty())
        {
            let (span, code) = line.split_once(": ").ok_or("No path")?;
            let span: CodeSpan = serde_json::from_str(span).map_err(|e| e.to_string())?;
            macro_spans.push(span.clone());
            let code: EcoString = serde_json::from_str(code).map_err(|e| e.to_string())?;
            macros.insert(span, code);
        }

        // Spans
        let mut spans = eco_vec![Span::Builtin];
        for line in spans_src.split('\n') {
            if line.trim().is_empty() {
                spans.push(Span::Builtin);
            } else {
                let (src_start, end) = line.trim().rsplit_once(' ').ok_or("invalid span")?;
                let (src, start) = src_start.split_once(' ').ok_or("invalid span")?;
                let src = if let Some(i) =
                    (src.strip_prefix("file")).and_then(|s| s.parse::<usize>().ok())
                {
                    InputSrc::File(PathBuf::from(&file_paths[i]).into())
                } else if let Some(i) =
                    (src.strip_prefix("macro")).and_then(|s| s.parse::<usize>().ok())
                {
                    InputSrc::Macro(macro_spans[i].clone().into())
                } else {
                    serde_json::from_str(src).map_err(|e| e.to_string())?
                };
                let start = serde_json::from_str(start).map_err(|e| e.to_string())?;
                let end = serde_json::from_str(end).map_err(|e| e.to_string())?;
                spans.push(Span::Code(CodeSpan { src, start, end }));
            }
        }
        assert_eq!(spans.pop(), Some(Span::Builtin)); // Remove extra empty line

        // Bindings
        let mut bindings = EcoVec::new();
        let mut lines = bindings_src
            .lines()
            .filter(|line| !line.trim().is_empty())
            .peekable();
        while let Some(line) = lines.next() {
            let (public, line) = if let Some(line) = line.strip_prefix("private ") {
                (false, line)
            } else {
                (true, line)
            };
            let (external, line) = if let Some(line) = line.strip_prefix("external ") {
                (true, line)
            } else {
                (false, line)
            };
            let (line, span) = line.rsplit_once(' ').ok_or("Missing binding span")?;
            let kind: BindingKind = serde_json::from_str(line).or_else(|e| {
                if let Some((key, val)) = line.split_once(' ') {
                    let json = format!("{{{key:?}: {val}}}");
                    serde_json::from_str(&json).map_err(|_| e.to_string())
                } else {
                    Err("No key".into())
                }
            })?;
            let span: usize = span
                .parse()
                .map_err(|_| format!("Invalid binding span {span:?}"))?;
            let span = spans
                .get(span)
                .ok_or_else(|| {
                    format!(
                        "Span index {} for binding {line:?} is out of bounds of {} spans",
                        span,
                        spans.len()
                    )
                })?
                .clone()
                .code()
                .unwrap_or_else(CodeSpan::dummy);
            let comment = (lines.peek())
                .and_then(|line| line.strip_prefix("  comment: "))
                .and_then(|s| {
                    lines.next();
                    serde_json::from_str::<DocComment>(s).ok()
                });
            let deprecation = (lines.peek())
                .and_then(|line| line.strip_prefix("  deprecation: "))
                .and_then(|s| {
                    lines.next();
                    serde_json::from_str::<EcoString>(s).ok()
                });
            bindings.push(BindingInfo {
                kind,
                public,
                span,
                meta: BindingMeta {
                    comment,
                    deprecation,
                    external,
                    ..Default::default()
                },
                used: true,
            });
        }

        // Functions
        let mut functions = EcoVec::new();
        for line in functions_src.lines().filter(|line| !line.trim().is_empty()) {
            let func: Node = serde_json::from_str(line).unwrap();
            functions.push(func);
        }

        // Index macros
        let mut index_macros = IndexMap::new();
        for line in index_macros_src
            .lines()
            .filter(|line| !line.trim().is_empty())
        {
            let (first, rest) = line.split_once(' ').ok_or("invalid index macro line")?;
            let i: usize = first
                .parse()
                .map_err(|_| format!("invalid index macro index {first:?}"))?;
            let index_macro: IndexMacro =
                serde_json::from_str(rest).map_err(|e| format!("invalid index macro #{i}: {e}"))?;
            index_macros.insert(i, index_macro);
        }
        let index_macros = Arc::new(index_macros);

        // Code macros
        let mut code_macros = IndexMap::new();
        for line in code_macros_src
            .lines()
            .filter(|line| !line.trim().is_empty())
        {
            let (first, rest) = line.split_once(' ').ok_or("invalid code macro line")?;
            let i: usize = first
                .parse()
                .map_err(|_| format!("invalid code macro code {first:?}"))?;
            let code_macro: CodeMacro =
                serde_json::from_str(rest).map_err(|e| format!("invalid code macro: {e}"))?;
            code_macros.insert(i, code_macro);
        }
        let code_macros = Arc::new(code_macros);

        // Strings
        let mut strings = EcoVec::new();
        for line in strings_src.lines() {
            let src: EcoString = serde_json::from_str(line).map_err(|e| e.to_string())?;
            strings.push(src);
        }

        Ok(Self {
            root,
            dependencies,
            exports,
            bindings,
            functions,
            index_macros,
            code_macros,
            spans,
            inputs: Inputs {
                files,
                strings,
                macros,
            },
            dynamic_functions: EcoVec::new(),
            test_assert_count: 0,
            line_sigs: BTreeMap::new(),
        })
    }
    /// Serialize the assembly into a `.uasm` file
    pub fn to_uasm(&self) -> String {
        let mut uasm = String::new();
        for node in self.root.iter() {
            uasm.push_str(&serde_json::to_string(node).unwrap());
            uasm.push('\n');
        }

        uasm.push_str("\nDEPENDENCIES\n");
        for (path, hash) in &self.dependencies {
            uasm.push_str(&format!("{} {hash}\n", path.display()));
        }

        uasm.push_str("\nEXPORTS\n");
        for (name, index) in &*self.exports {
            uasm.push_str(&format!("{name} {index}\n"));
        }

        uasm.push_str("\nBINDINGS\n");
        let span_indices: HashMap<&CodeSpan, usize> = (self.spans.iter().enumerate())
            .filter_map(|(i, span)| span.code_ref().map(|s| (s, i)))
            .collect();
        for binding in &self.bindings {
            if !binding.public {
                uasm.push_str("private ");
            }
            if binding.meta.external {
                uasm.push_str("external ");
            }
            if let serde_json::Value::Object(map) = serde_json::to_value(&binding.kind).unwrap()
                && map.len() == 1
            {
                let key = map.keys().next().unwrap();
                let value = map.values().next().unwrap();
                uasm.push_str(&format!("{key} {value}"));
            } else {
                uasm.push_str(&serde_json::to_string(&binding.kind).unwrap());
            }
            uasm.push(' ');
            let span = (span_indices.get(&binding.span).copied().unwrap_or(0)).to_string();
            uasm.push_str(&span);
            uasm.push('\n');
            if let Some(com) = &binding.meta.comment {
                uasm.push_str("  comment: ");
                uasm.push_str(&serde_json::to_string(com).unwrap());
                uasm.push('\n');
            }
            if let Some(deprecation) = &binding.meta.deprecation {
                uasm.push_str("  deprecation: ");
                uasm.push_str(&serde_json::to_string(deprecation).unwrap());
                uasm.push('\n');
            }
        }

        uasm.push_str("\nFUNCTIONS\n");
        for func in &self.functions {
            uasm.push_str(&serde_json::to_string(&func).unwrap());
            uasm.push('\n');
        }

        uasm.push_str("\nINDEX MACROS\n");
        for (i, mac) in &*self.index_macros {
            uasm.push_str(&format!("{i} {}", serde_json::to_string(mac).unwrap()));
            uasm.push('\n');
        }

        uasm.push_str("\nCODE MACROS\n");
        for (i, mac) in &*self.code_macros {
            uasm.push_str(&format!("{i} {}", serde_json::to_string(mac).unwrap()));
            uasm.push('\n');
        }

        uasm.push_str("\nSPANS\n");
        let file_indices: HashMap<PathBuf, usize> = (self.inputs.files.iter().enumerate())
            .map(|(i, entry)| (entry.key().into(), i))
            .collect();
        let macro_indices: HashMap<CodeSpan, usize> = (self.inputs.macros.iter().enumerate())
            .map(|(i, entry)| (entry.key().clone(), i))
            .collect();
        for span in self.spans.iter().skip(1) {
            if let Span::Code(span) = span {
                let src = match &span.src {
                    InputSrc::File(path) => format!("file{}", file_indices[&**path]),
                    InputSrc::Macro(span) => format!("macro{}", macro_indices[span]),
                    src => serde_json::to_string(src).unwrap(),
                };
                uasm.push_str(&src);
                uasm.push(' ');
                uasm.push_str(&serde_json::to_string(&span.start).unwrap());
                uasm.push(' ');
                uasm.push_str(&serde_json::to_string(&span.end).unwrap());
            }
            uasm.push('\n');
        }

        uasm.push_str("\nFILES\n");
        for entry in &self.inputs.files {
            uasm.push_str(&format!("{}: {:?}\n", entry.key().display(), entry.value()));
        }

        uasm.push_str("\nMACRO EXPANSIONS\n");
        for entry in &self.inputs.macros {
            uasm.push_str(&serde_json::to_string(entry.key()).unwrap());
            uasm.push_str(&format!(": {:?}\n", entry.value()));
        }

        if !self.inputs.strings.is_empty() {
            uasm.push_str("\nSTRING INPUTS\n");
            for src in &self.inputs.strings {
                uasm.push_str(&serde_json::to_string(src).unwrap());
                uasm.push('\n');
            }
        }

        uasm
    }
}

impl Index<&Function> for Assembly {
    type Output = Node;
    #[track_caller]
    fn index(&self, func: &Function) -> &Self::Output {
        match self.functions.get(func.index) {
            Some(node) => node,
            None => panic!("{}({:?}) not found in assembly", func.id, func.index),
        }
    }
}

impl IndexMut<&Function> for Assembly {
    #[track_caller]
    fn index_mut(&mut self, func: &Function) -> &mut Self::Output {
        match self.functions.make_mut().get_mut(func.index) {
            Some(node) => node,
            None => panic!("{}({:?}) not found in assembly", func.id, func.index),
        }
    }
}

#[cfg(not(target_arch = "wasm32"))]
type DynFn = Arc<dyn Fn(&mut Uiua) -> UiuaResult + Send + Sync + 'static>;
#[cfg(target_arch = "wasm32")]
type DynFn = Arc<dyn Fn(&mut Uiua) -> UiuaResult + 'static>;

impl Default for Assembly {
    fn default() -> Self {
        Self {
            root: Node::default(),
            dependencies: EcoVec::new(),
            exports: Arc::new(IndexMap::new()),
            functions: EcoVec::new(),
            index_macros: Arc::new(IndexMap::new()),
            code_macros: Arc::new(IndexMap::new()),
            spans: eco_vec![Span::Builtin],
            bindings: EcoVec::new(),
            dynamic_functions: EcoVec::new(),
            inputs: Inputs::default(),
            test_assert_count: 0,
            line_sigs: BTreeMap::new(),
        }
    }
}

impl From<&Assembly> for Assembly {
    fn from(asm: &Assembly) -> Self {
        asm.clone()
    }
}

/// Information about a binding
#[derive(Debug, Clone)]
pub struct BindingInfo {
    /// The binding kind
    pub kind: BindingKind,
    /// Whether the binding is public
    pub public: bool,
    /// The span of the original binding name
    pub span: CodeSpan,
    /// Metadata about the binding
    pub meta: BindingMeta,
    /// Whether the binding was used
    pub used: bool,
}

/// Metadata about a binding
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct BindingMeta {
    /// The comment preceding the binding
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub comment: Option<DocComment>,
    /// The character counts for golfing
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub counts: Option<BindingCounts>,
    /// The deprecation message
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub deprecation: Option<EcoString>,
    /// Whether this binding's code was externally provided
    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
    pub external: bool,
}

/// A kind of global binding
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum BindingKind {
    /// A constant value
    Const(Option<Value>),
    /// A function
    Func(Function),
    /// A module
    Module(Module),
    /// A scope being compiled
    Scope(usize),
    /// An index macro
    ///
    /// Contains the number of arguments
    IndexMacro(usize),
    /// A code macro
    CodeMacro(Node),
    /// An error
    Error,
}

impl BindingKind {
    /// Get the signature of the binding
    pub fn sig(&self) -> Option<Signature> {
        match self {
            Self::Const(_) => Some(Signature::new(0, 1)),
            Self::Func(func) => Some(func.sig),
            Self::Module(_) => None,
            Self::Scope(_) => None,
            Self::IndexMacro(_) => None,
            Self::CodeMacro(_) => None,
            Self::Error => None,
        }
    }
    /// Check if the binding is a once-bound constant
    pub fn is_constant(&self) -> bool {
        matches!(self, Self::Const(_))
    }
    /// Check if the binding is a module
    pub fn is_module(&self) -> bool {
        self.as_module().is_some()
    }
    /// Get the binding as a module
    pub fn as_module(&self) -> Option<&Module> {
        match self {
            BindingKind::Module(m) => Some(m),
            _ => None,
        }
    }
    /// Check if the binding is a constant or function
    pub fn has_sig(&self) -> bool {
        match self {
            Self::Const(_) | Self::Func(_) => true,
            Self::Module(m) => m.names.contains_key("Call") || m.names.contains_key("New"),
            _ => false,
        }
    }
}

/// A comment that documents a binding
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct DocComment {
    /// The comment text
    #[serde(default, skip_serializing_if = "str::is_empty")]
    pub text: EcoString,
    /// The signature of the binding
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub sig: Option<DocCommentSig>,
}

/// A signature in a doc comment
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct DocCommentSig {
    /// Whether this is a labelling signature
    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
    pub label: bool,
    /// The arguments of the signature
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub args: Option<Vec<DocCommentArg>>,
    /// The outputs of the signature
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub outputs: Option<Vec<DocCommentArg>>,
}

impl DocCommentSig {
    /// Whether the doc comment signature matches a given function signature
    pub fn matches_sig(&self, sig: Signature) -> bool {
        (self.args.as_ref()).is_none_or(|args| args.len() == sig.args())
            && (self.outputs.as_ref()).is_none_or(|o| o.len() == sig.outputs())
    }
    pub(crate) fn sig_string(&self) -> String {
        match (&self.args, &self.outputs) {
            (Some(args), Some(outputs)) => {
                format!("signature {}", Signature::new(args.len(), outputs.len()))
            }
            (Some(args), None) => format!(
                "{} arg{}",
                args.len(),
                if args.len() == 1 { "" } else { "s" }
            ),
            (None, Some(outputs)) => format!(
                "{} output{}",
                outputs.len(),
                if outputs.len() == 1 { "" } else { "s" }
            ),
            (None, None) => "signature".into(),
        }
    }
}

impl fmt::Display for DocCommentSig {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if let Some(outputs) = &self.outputs {
            for output in outputs {
                write!(f, " {}", output.name)?;
                if let Some(ty) = &output.ty {
                    write!(f, ":{ty}")?;
                }
            }
            write!(f, " ")?;
        }
        if self.label {
            write!(f, "$")?;
        } else {
            write!(f, "?")?;
        }
        if let Some(args) = &self.args {
            write!(f, " ")?;
            for (i, arg) in args.iter().enumerate() {
                if i > 0 {
                    write!(f, " ")?;
                }
                write!(f, "{}", arg.name)?;
                if let Some(ty) = &arg.ty {
                    write!(f, ":{ty}")?;
                }
            }
        }
        Ok(())
    }
}

/// An argument in a doc comment signature
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct DocCommentArg {
    /// The name of the argument
    pub name: EcoString,
    /// A type descriptor for the argument
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub ty: Option<EcoString>,
}

fn is_sig_line(s: &str) -> bool {
    if s.chars().filter(|&c| "$?".contains(c)).count() != 1 {
        return false;
    }
    let s = s.trim_end();
    (!s.ends_with(['?', '$']) || s.ends_with(" ?") || s.ends_with(" $"))
        && (s.chars()).all(|c| {
            c.is_whitespace()
                || "?$:".contains(c)
                || is_ident_char(c)
                || SUBSCRIPT_DIGITS.contains(&c)
        })
}

impl FromStr for DocCommentSig {
    type Err = ();
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if !is_sig_line(s) {
            return Err(());
        }
        // Split into args and outputs
        let mut label = false;
        let (mut outputs_text, mut args_text) = s
            .split_once('?')
            .or_else(|| s.split_once('$').inspect(|_| label = true))
            .ok_or(())?;
        outputs_text = outputs_text.trim();
        args_text = args_text.trim();
        // Parse args and outputs
        let mut args = Vec::new();
        let mut outputs = Vec::new();
        for (args, text) in [(&mut args, args_text), (&mut outputs, outputs_text)] {
            // Tokenize text
            let mut tokens = Vec::new();
            for frag in text.split_whitespace() {
                for (i, token) in frag.split(':').enumerate() {
                    if i > 0 {
                        tokens.push(":");
                    }
                    tokens.push(token);
                }
            }
            // Parse tokens into args
            let mut curr_arg_name = None;
            let mut tokens = tokens.into_iter().peekable();
            while let Some(token) = tokens.next() {
                if token == ":" {
                    let ty = tokens.next().unwrap_or_default();
                    args.push(DocCommentArg {
                        name: curr_arg_name.take().unwrap_or_default(),
                        ty: if ty.is_empty() { None } else { Some(ty.into()) },
                    });
                } else {
                    if let Some(curr) = curr_arg_name.take() {
                        args.push(DocCommentArg {
                            name: curr,
                            ty: None,
                        });
                    }
                    curr_arg_name = Some(token.into());
                }
            }
            if let Some(curr) = curr_arg_name.take() {
                args.push(DocCommentArg {
                    name: curr,
                    ty: None,
                });
            }
        }
        Ok(DocCommentSig {
            label,
            args: (!args.is_empty()).then_some(args),
            outputs: (!outputs.is_empty()).then_some(outputs),
        })
    }
}

impl From<String> for DocComment {
    fn from(text: String) -> Self {
        Self::from(text.as_str())
    }
}

impl From<&str> for DocComment {
    fn from(text: &str) -> Self {
        let mut sig = None;
        let sig_line = text.lines().position(is_sig_line);
        let raw_text = if let Some(i) = sig_line {
            sig = text.lines().nth(i).unwrap().parse().ok();

            let mut text: EcoString = (text.lines().take(i))
                .chain(["\n"])
                .chain(text.lines().skip(i + 1))
                .flat_map(|s| s.chars().chain(Some('\n')))
                .collect();
            while text.ends_with('\n') {
                text.pop();
            }
            if text.starts_with('\n') {
                text = text.trim_start_matches('\n').into();
            }
            text
        } else {
            text.into()
        };
        let mut text = EcoString::new();
        for (i, line) in raw_text.lines().enumerate() {
            if i > 0 {
                text.push('\n');
            }
            text.push_str(line.trim());
        }
        DocComment { text, sig }
    }
}

impl fmt::Debug for Assembly {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        struct FmtFunctions<'a>(&'a Assembly);
        impl fmt::Debug for FmtFunctions<'_> {
            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                f.debug_list()
                    .entries(self.0.bindings.iter().filter_map(|b| {
                        if let BindingKind::Func(func) = &b.kind {
                            Some((func, &self.0[func]))
                        } else {
                            None
                        }
                    }))
                    .finish()
            }
        }
        f.debug_struct("Assembly")
            .field("root", &self.root)
            .field("functions", &FmtFunctions(self))
            .finish()
    }
}