flowlog-build 0.2.2

//! Atom types for FlowLog Datalog programs.
//!
//! - [`AtomArg`]: variable / constant / placeholder (`_`)
//! - [`Atom`]: `name(arg1, ..., argN)`

use crate::parser::error::{grammar_bug, ParseError};
use crate::parser::primitive::ConstType;
use crate::parser::{span_of, Lexeme, Rule};
use pest::iterators::Pair;
use std::fmt;

use crate::common::compute_fp;
use crate::common::{FileId, Ignored, Span};

/// An argument to an atom: variable, constant, or `_`.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub(crate) enum AtomArg {
    Var(String),
    Const(ConstType),
    Placeholder,
}

impl fmt::Display for AtomArg {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Var(v) => write!(f, "{v}"),
            Self::Const(c) => write!(f, "{c}"),
            Self::Placeholder => write!(f, "_"),
        }
    }
}

impl Lexeme for AtomArg {
    /// Parse an atom argument from the grammar.
    fn from_parsed_rule(parsed_rule: Pair<Rule>, file: FileId) -> Result<Self, ParseError> {
        let inner = parsed_rule
            .into_inner()
            .next()
            .ok_or_else(|| grammar_bug("atom_arg missing inner token"))?;

        Ok(match inner.as_rule() {
            Rule::variable => Self::Var(inner.as_str().to_string()),
            Rule::constant => Self::Const(ConstType::from_parsed_rule(inner, file)?),
            Rule::placeholder => Self::Placeholder,
            other => {
                return Err(grammar_bug(format!(
                    "invalid atom argument rule: {other:?}"
                )))
            }
        })
    }
}

/// `name(arg1, ..., argN)` predicate.
#[derive(Clone, PartialEq, Eq, Hash)]
pub(crate) struct Atom {
    name: String,
    arguments: Vec<AtomArg>,
    fingerprint: u64,
    span: Ignored<Span>,
}

impl Atom {
    /// Create a new atom.
    ///
    /// Converts the name to lowercase.
    #[must_use]
    pub(crate) fn new(name: &str, arguments: Vec<AtomArg>, fingerprint: u64) -> Self {
        Self {
            name: name.to_lowercase(),
            arguments,
            fingerprint,
            span: Ignored(Span::DUMMY),
        }
    }

    /// Source location this atom was parsed from.
    #[must_use]
    #[inline]
    pub(crate) fn span(&self) -> Span {
        self.span.0
    }

    /// Relation name.
    #[must_use]
    pub(crate) fn name(&self) -> &str {
        &self.name
    }

    /// Arguments (as a slice).
    #[must_use]
    pub(crate) fn arguments(&self) -> &[AtomArg] {
        &self.arguments
    }

    pub(crate) fn arguments_mut(&mut self) -> &mut [AtomArg] {
        &mut self.arguments
    }

    /// Number of arguments.
    #[must_use]
    pub(crate) fn arity(&self) -> usize {
        self.arguments.len()
    }

    /// Get the fingerprint.
    #[must_use]
    pub(crate) fn fingerprint(&self) -> u64 {
        self.fingerprint
    }
}

impl fmt::Display for Atom {
    /// Formats as `name(a, b, _)`, always including parentheses.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}(", self.name)?;
        for (i, arg) in self.arguments.iter().enumerate() {
            if i > 0 {
                write!(f, ", ")?;
            }
            write!(f, "{arg}")?;
        }
        write!(f, ")")
    }
}

impl fmt::Debug for Atom {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}(", self.name)?;
        for (i, arg) in self.arguments.iter().enumerate() {
            if i > 0 {
                write!(f, ", ")?;
            }
            write!(f, "{arg}")?;
        }
        write!(f, ") [fp: 0x{:016x}]", self.fingerprint)
    }
}

impl Lexeme for Atom {
    /// Parse `name("(" (atom_arg ("," atom_arg)*)? ")")`.
    fn from_parsed_rule(parsed_rule: Pair<Rule>, file: FileId) -> Result<Self, ParseError> {
        let span = span_of(&parsed_rule, file);
        let mut inner = parsed_rule.into_inner();

        let name = inner
            .next()
            .ok_or_else(|| grammar_bug("atom missing relation name"))?
            .as_str()
            .to_lowercase();
        let fingerprint = compute_fp(&name);

        let mut arguments = Vec::new();
        for pair in inner {
            if pair.as_rule() == Rule::atom_arg {
                arguments.push(AtomArg::from_parsed_rule(pair, file)?);
            }
        }

        Ok(Self {
            name,
            arguments,
            fingerprint,
            span: Ignored(span),
        })
    }
}