vamp/src/snippet.rs

use std::{collections::BTreeMap, ops::Range};

use unicode_segmentation::UnicodeSegmentation;

use crate::plugin::{CompletionSnippet, SnippetTabstop};

const MAX_BYTES: usize = 256 * 1024;
const MAX_NODES: usize = 4096;
const MAX_STOPS: usize = 1024;
const MAX_DEPTH: usize = 32;
const MAX_WORK: usize = 1024 * 1024;

#[derive(Debug, Eq, PartialEq)]
pub(crate) enum Error {
    Invalid,
    Unsupported,
    Limit,
}

enum Node<'a> {
    Text(&'a str),
    Value(String),
    Stop(u32, Option<Vec<usize>>),
    Default(Vec<usize>),
}

struct Parser<'a, F> {
    text: &'a str,
    cursor: usize,
    nodes: Vec<Node<'a>>,
    value_bytes: usize,
    variable: F,
}

/// Parse normalized snippet text before assigning any byte offsets. Unsupported
/// constructs omit the entire candidate, never leaking template syntax into text.
pub(crate) fn expand(
    text: &str,
    variable: impl Fn(&str) -> Option<String>,
) -> Result<(String, CompletionSnippet), Error> {
    if text.len() > MAX_BYTES {
        return Err(Error::Limit);
    }
    let normalized = if text.contains('\r') {
        std::borrow::Cow::Owned(text.replace("\r\n", "\n").replace('\r', "\n"))
    } else {
        std::borrow::Cow::Borrowed(text)
    };
    let mut parser = Parser {
        text: &normalized,
        cursor: 0,
        nodes: Vec::new(),
        value_bytes: 0,
        variable,
    };
    let roots = parser.sequence(false, 0)?;
    let mut definitions = BTreeMap::new();
    let mut pending: Vec<_> = roots.iter().rev().copied().collect();
    while let Some(id) = pending.pop() {
        match &parser.nodes[id] {
            Node::Stop(index, Some(children)) => {
                definitions.entry(*index).or_insert(id);
                pending.extend(children.iter().rev().copied());
            }
            Node::Default(children) => pending.extend(children.iter().rev().copied()),
            _ => {}
        }
    }
    let mut rendered = Render {
        nodes: &parser.nodes,
        definitions: &definitions,
        text: String::new(),
        stops: BTreeMap::new(),
        ancestors: Vec::new(),
        occurrences: 0,
        work: 0,
    };
    rendered.sequence(&roots, 0)?;
    if !rendered.stops.contains_key(&0) {
        if rendered.occurrences == MAX_STOPS {
            return Err(Error::Limit);
        }
        rendered.stops.insert(
            0,
            SnippetTabstop {
                index: 0,
                ranges: std::iter::once(rendered.text.len()..rendered.text.len()).collect(),
                ancestors: vec![Vec::new()],
            },
        );
    }
    let final_stop = rendered.stops.get(&0).unwrap();
    if final_stop.ranges.len() != 1
        || !final_stop.ranges[0].is_empty()
        || !final_stop.ancestors[0].is_empty()
        || rendered.stops.values().any(|stop| {
            stop.ancestors
                .iter()
                .any(|ancestors| ancestors.contains(&0))
        })
    {
        return Err(Error::Unsupported);
    }
    let boundaries: Vec<_> = rendered
        .text
        .grapheme_indices(true)
        .map(|(offset, _)| offset)
        .chain(std::iter::once(rendered.text.len()))
        .collect();
    for stop in rendered.stops.values() {
        let mut previous: Option<&Range<usize>> = None;
        for range in &stop.ranges {
            if boundaries.binary_search(&range.start).is_err()
                || boundaries.binary_search(&range.end).is_err()
                || previous.is_some_and(|old| {
                    old.end > range.start
                        || (old.start == range.start && !(old.is_empty() && range.is_empty()))
                })
            {
                return Err(Error::Invalid);
            }
            previous = Some(range);
        }
    }
    let final_stop = rendered.stops.remove(&0).unwrap();
    let mut tabstops: Vec<_> = rendered.stops.into_values().collect();
    tabstops.push(final_stop);
    Ok((rendered.text, CompletionSnippet { tabstops }))
}

impl<'a, F: Fn(&str) -> Option<String>> Parser<'a, F> {
    fn node(&mut self, mut node: Node<'a>) -> Result<usize, Error> {
        if self.nodes.len() == MAX_NODES {
            return Err(Error::Limit);
        }
        if let Node::Value(value) = &mut node {
            if value.len() > MAX_BYTES {
                return Err(Error::Limit);
            }
            if value.contains('\r') {
                *value = value.replace("\r\n", "\n").replace('\r', "\n");
            }
            self.value_bytes = self
                .value_bytes
                .checked_add(value.len())
                .filter(|size| *size <= MAX_BYTES)
                .ok_or(Error::Limit)?;
        }
        let id = self.nodes.len();
        self.nodes.push(node);
        Ok(id)
    }

    fn byte(&self) -> Option<u8> {
        self.text.as_bytes().get(self.cursor).copied()
    }

    fn sequence(&mut self, nested: bool, depth: usize) -> Result<Vec<usize>, Error> {
        if depth > MAX_DEPTH {
            return Err(Error::Limit);
        }
        let mut nodes = Vec::new();
        while let Some(byte) = self.byte() {
            if nested && byte == b'}' {
                self.cursor += 1;
                return Ok(nodes);
            }
            if byte == b'$' {
                nodes.push(self.construct(depth)?);
            } else if byte == b'\\' {
                self.cursor += 1;
                let start = self.cursor;
                if !matches!(self.byte(), Some(b'$' | b'}' | b'\\')) {
                    return Err(Error::Invalid);
                }
                self.cursor += 1;
                nodes.push(self.node(Node::Text(&self.text[start..self.cursor]))?);
            } else {
                let start = self.cursor;
                while let Some(next) = self.byte() {
                    if next == b'$' || next == b'\\' || (nested && next == b'}') {
                        break;
                    }
                    self.cursor += 1;
                }
                nodes.push(self.node(Node::Text(&self.text[start..self.cursor]))?);
            }
        }
        if nested {
            Err(Error::Invalid)
        } else {
            Ok(nodes)
        }
    }

    fn construct(&mut self, depth: usize) -> Result<usize, Error> {
        self.cursor += 1;
        let braced = self.byte() == Some(b'{');
        self.cursor += usize::from(braced);
        let start = self.cursor;
        let numbered = self.byte().is_some_and(|byte| byte.is_ascii_digit());
        if numbered {
            while self.byte().is_some_and(|byte| byte.is_ascii_digit()) {
                self.cursor += 1;
            }
        } else {
            if !self
                .byte()
                .is_some_and(|byte| byte.is_ascii_alphabetic() || byte == b'_')
            {
                return Err(Error::Invalid);
            }
            self.cursor += 1;
            while self
                .byte()
                .is_some_and(|byte| byte.is_ascii_alphanumeric() || byte == b'_')
            {
                self.cursor += 1;
            }
        }
        let name = &self.text[start..self.cursor];
        let mut children = None;
        if braced {
            match self.byte() {
                Some(b'}') => self.cursor += 1,
                Some(b':') => {
                    self.cursor += 1;
                    children = Some(self.sequence(true, depth + 1)?);
                }
                Some(b'|') if numbered => {
                    self.cursor += 1;
                    children = Some(self.choice()?);
                }
                Some(b'/') => return Err(Error::Unsupported),
                _ => return Err(Error::Invalid),
            }
        }
        if numbered {
            let index = name.parse::<u32>().map_err(|_| Error::Limit)?;
            self.node(Node::Stop(index, children))
        } else {
            match (self.variable)(name) {
                Some(value) if !value.is_empty() => self.node(Node::Value(value)),
                value => match children {
                    Some(children) => self.node(Node::Default(children)),
                    None if value.is_some() => self.node(Node::Text("")),
                    None => Err(Error::Unsupported),
                },
            }
        }
    }

    fn choice(&mut self) -> Result<Vec<usize>, Error> {
        let mut first = String::new();
        let mut keep = true;
        loop {
            match self.byte().ok_or(Error::Invalid)? {
                b'|' => {
                    self.cursor += 1;
                    if self.byte() != Some(b'}') {
                        return Err(Error::Invalid);
                    }
                    self.cursor += 1;
                    return Ok(vec![self.node(Node::Value(first))?]);
                }
                b',' => {
                    keep = false;
                    self.cursor += 1;
                }
                b'\\' => {
                    self.cursor += 1;
                    let byte = self.byte().ok_or(Error::Invalid)?;
                    if !matches!(byte, b'$' | b'}' | b'\\' | b',' | b'|') {
                        return Err(Error::Invalid);
                    }
                    if keep {
                        first.push(byte as char);
                    }
                    self.cursor += 1;
                }
                _ => {
                    let ch = self.text[self.cursor..].chars().next().unwrap();
                    if keep {
                        first.push(ch);
                    }
                    self.cursor += ch.len_utf8();
                }
            }
        }
    }
}

struct Render<'a> {
    nodes: &'a [Node<'a>],
    definitions: &'a BTreeMap<u32, usize>,
    text: String,
    stops: BTreeMap<u32, SnippetTabstop>,
    ancestors: Vec<u32>,
    occurrences: usize,
    work: usize,
}

impl Render<'_> {
    fn sequence(&mut self, nodes: &[usize], depth: usize) -> Result<(), Error> {
        if depth > MAX_DEPTH {
            return Err(Error::Limit);
        }
        let arena = self.nodes;
        for &id in nodes {
            self.work += 1;
            if self.work > MAX_WORK {
                return Err(Error::Limit);
            }
            match &arena[id] {
                Node::Text(text) => self.append(text)?,
                Node::Value(text) => self.append(text)?,
                Node::Default(children) => self.sequence(children, depth + 1)?,
                Node::Stop(index, _) => {
                    if self.occurrences == MAX_STOPS {
                        return Err(Error::Limit);
                    }
                    self.occurrences += 1;
                    let start = self.text.len();
                    let ancestors = self.ancestors.clone();
                    if let Some(&definition) = self.definitions.get(index) {
                        if self.ancestors.contains(index) {
                            return Err(Error::Invalid);
                        }
                        self.ancestors.push(*index);
                        if let Node::Stop(_, Some(children)) = &arena[definition] {
                            self.sequence(children, depth + 1)?;
                        }
                        self.ancestors.pop();
                    }
                    let stop = self.stops.entry(*index).or_insert_with(|| SnippetTabstop {
                        index: *index,
                        ranges: Vec::new(),
                        ancestors: Vec::new(),
                    });
                    stop.ranges.push(start..self.text.len());
                    stop.ancestors.push(ancestors);
                }
            }
        }
        Ok(())
    }

    fn append(&mut self, text: &str) -> Result<(), Error> {
        if self
            .text
            .len()
            .checked_add(text.len())
            .is_none_or(|size| size > MAX_BYTES)
        {
            return Err(Error::Limit);
        }
        self.text.push_str(text);
        Ok(())
    }
}

#[cfg(test)]
// These fixtures intentionally contain ranges, not lists of indices.
#[allow(clippy::single_range_in_vec_init)]
#[path = "../tests/unit/snippet.rs"]
mod tests;