tensorlogic_ir/

refinement.rs

//! Refinement types for constraint-based type checking.
//!
//! Refinement types extend base types with logical predicates that constrain
//! the valid values of that type. This enables more precise type checking and
//! verification.
//!
//! # Examples
//!
//! ```
//! use tensorlogic_ir::refinement::{RefinementType, Refinement};
//! use tensorlogic_ir::TLExpr;
//!
//! // Positive integers: {x: Int | x > 0}
//! let positive_int = RefinementType::new(
//!     "x",
//!     "Int",
//!     TLExpr::gt(TLExpr::pred("x", vec![]), TLExpr::constant(0.0))
//! );
//!
//! // Bounded values: {x: Float | x >= 0.0 && x <= 1.0}
//! let probability = RefinementType::new(
//!     "x",
//!     "Float",
//!     TLExpr::and(
//!         TLExpr::gte(TLExpr::pred("x", vec![]), TLExpr::constant(0.0)),
//!         TLExpr::lte(TLExpr::pred("x", vec![]), TLExpr::constant(1.0))
//!     )
//! );
//! ```

use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::fmt;

use crate::{IrError, ParametricType, TLExpr, Term};

/// Refinement: a logical predicate that refines a type.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Refinement {
    /// Variable name being refined
    pub var_name: String,
    /// Refinement predicate
    pub predicate: TLExpr,
}

impl Refinement {
    pub fn new(var_name: impl Into<String>, predicate: TLExpr) -> Self {
        Refinement {
            var_name: var_name.into(),
            predicate,
        }
    }

    /// Get free variables in the refinement (excluding the refined variable)
    pub fn free_vars(&self) -> HashSet<String> {
        let mut vars = self.predicate.free_vars();
        vars.remove(&self.var_name);
        vars
    }

    /// Substitute variables in the refinement
    pub fn substitute(&self, subst: &HashMap<String, Term>) -> Refinement {
        // Don't substitute the refined variable itself
        let mut filtered_subst = subst.clone();
        filtered_subst.remove(&self.var_name);

        Refinement {
            var_name: self.var_name.clone(),
            predicate: self.predicate.clone(), // Would need substitute method on TLExpr
        }
    }

    /// Simplify the refinement predicate
    pub fn simplify(&self) -> Refinement {
        use crate::optimize_expr;

        Refinement {
            var_name: self.var_name.clone(),
            predicate: optimize_expr(&self.predicate),
        }
    }

    /// Check if refinement implies another refinement
    pub fn implies(&self, other: &Refinement) -> bool {
        // Simplified check - would need SMT solver for full verification
        if self.var_name != other.var_name {
            return false;
        }

        // Syntactic equality check
        self.predicate == other.predicate
    }
}

impl fmt::Display for Refinement {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{{{}: | {}}}", self.var_name, self.predicate)
    }
}

/// Refinement type: base type with a refinement predicate.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct RefinementType {
    /// Variable name
    pub var_name: String,
    /// Base type
    pub base_type: ParametricType,
    /// Refinement predicate on the variable
    pub refinement: TLExpr,
}

impl RefinementType {
    pub fn new(
        var_name: impl Into<String>,
        base_type: impl Into<String>,
        refinement: TLExpr,
    ) -> Self {
        RefinementType {
            var_name: var_name.into(),
            base_type: ParametricType::concrete(base_type),
            refinement,
        }
    }

    /// Create a refinement type from parametric type
    pub fn from_parametric(
        var_name: impl Into<String>,
        base_type: ParametricType,
        refinement: TLExpr,
    ) -> Self {
        RefinementType {
            var_name: var_name.into(),
            base_type,
            refinement,
        }
    }

    /// Positive integers: {x: Int | x > 0}
    pub fn positive_int(var_name: impl Into<String>) -> Self {
        let var_name = var_name.into();
        RefinementType::new(
            var_name.clone(),
            "Int",
            TLExpr::gt(TLExpr::pred(&var_name, vec![]), TLExpr::constant(0.0)),
        )
    }

    /// Non-negative integers: {x: Int | x >= 0}
    pub fn nat(var_name: impl Into<String>) -> Self {
        let var_name = var_name.into();
        RefinementType::new(
            var_name.clone(),
            "Int",
            TLExpr::gte(TLExpr::pred(&var_name, vec![]), TLExpr::constant(0.0)),
        )
    }

    /// Probability: {x: Float | x >= 0.0 && x <= 1.0}
    pub fn probability(var_name: impl Into<String>) -> Self {
        let var_name = var_name.into();
        RefinementType::new(
            var_name.clone(),
            "Float",
            TLExpr::and(
                TLExpr::gte(TLExpr::pred(&var_name, vec![]), TLExpr::constant(0.0)),
                TLExpr::lte(TLExpr::pred(&var_name, vec![]), TLExpr::constant(1.0)),
            ),
        )
    }

    /// Non-empty vector: `{v: Vec<T> | length(v) > 0}`
    pub fn non_empty_vec(var_name: impl Into<String>, element_type: impl Into<String>) -> Self {
        let var_name = var_name.into();
        use crate::TypeConstructor;

        let elem_type = ParametricType::concrete(element_type);
        let vec_type = ParametricType::apply(TypeConstructor::List, vec![elem_type]);

        RefinementType::from_parametric(
            var_name.clone(),
            vec_type,
            TLExpr::gt(TLExpr::pred("length", vec![]), TLExpr::constant(0.0)),
        )
    }

    /// Get free variables in the refinement (excluding the refined variable)
    pub fn free_vars(&self) -> HashSet<String> {
        let mut vars = self.refinement.free_vars();
        vars.remove(&self.var_name);
        vars
    }

    /// Check if this type is a subtype of another
    pub fn is_subtype_of(&self, other: &RefinementType) -> bool {
        // Base types must match
        if self.base_type != other.base_type {
            return false;
        }

        // Refined variables must match
        if self.var_name != other.var_name {
            return false;
        }

        // self's refinement must imply other's refinement
        // (would need SMT solver for full verification)
        self.refinement == other.refinement
    }

    /// Weaken the refinement (make it less restrictive)
    pub fn weaken(&self) -> RefinementType {
        // Remove the refinement, keeping only the base type
        RefinementType {
            var_name: self.var_name.clone(),
            base_type: self.base_type.clone(),
            refinement: TLExpr::constant(1.0), // Always true
        }
    }

    /// Strengthen the refinement (add more constraints)
    pub fn strengthen(&self, additional: TLExpr) -> RefinementType {
        RefinementType {
            var_name: self.var_name.clone(),
            base_type: self.base_type.clone(),
            refinement: TLExpr::and(self.refinement.clone(), additional),
        }
    }
}

impl fmt::Display for RefinementType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{{{}: {} | {}}}",
            self.var_name, self.base_type, self.refinement
        )
    }
}

/// Refinement type checking context.
#[derive(Clone, Debug, Default)]
pub struct RefinementContext {
    /// Type bindings
    bindings: HashMap<String, RefinementType>,
    /// Assumed facts (refinement predicates that are known to be true)
    assumptions: Vec<TLExpr>,
}

impl RefinementContext {
    pub fn new() -> Self {
        Self::default()
    }

    /// Bind a variable to a refinement type
    pub fn bind(&mut self, name: impl Into<String>, typ: RefinementType) {
        let name = name.into();

        // Add the refinement as an assumption with variable substitution
        let assumption = typ.refinement.clone();
        self.assumptions.push(assumption);

        self.bindings.insert(name, typ);
    }

    /// Get the type of a variable
    pub fn get_type(&self, name: &str) -> Option<&RefinementType> {
        self.bindings.get(name)
    }

    /// Add an assumption
    pub fn assume(&mut self, fact: TLExpr) {
        self.assumptions.push(fact);
    }

    /// Check if a refinement is satisfied under current assumptions
    pub fn check_refinement(&self, refinement: &TLExpr) -> bool {
        // Simplified check - would need SMT solver for full verification
        // For now, check if the refinement is in our assumptions
        self.assumptions.contains(refinement)
    }

    /// Verify that a value satisfies a refinement type
    pub fn verify(&self, _value: &Term, _typ: &RefinementType) -> Result<(), IrError> {
        // Would need symbolic execution or SMT solving
        // For now, just check that the refinement is satisfiable
        Ok(())
    }
}

/// Liquid types: refinement types with inference.
#[derive(Clone, Debug)]
pub struct LiquidTypeInference {
    context: RefinementContext,
    /// Unknown refinements to be inferred
    unknowns: HashMap<String, Vec<TLExpr>>,
}

impl LiquidTypeInference {
    pub fn new() -> Self {
        LiquidTypeInference {
            context: RefinementContext::new(),
            unknowns: HashMap::new(),
        }
    }

    /// Add an unknown refinement variable
    pub fn add_unknown(&mut self, name: impl Into<String>, candidates: Vec<TLExpr>) {
        self.unknowns.insert(name.into(), candidates);
    }

    /// Infer refinements based on constraints
    pub fn infer(&mut self) -> HashMap<String, TLExpr> {
        // Simplified inference - would need constraint solving
        let mut inferred = HashMap::new();

        for (name, candidates) in &self.unknowns {
            // Pick the weakest (least restrictive) candidate that is satisfiable
            if let Some(refinement) = candidates.first() {
                inferred.insert(name.clone(), refinement.clone());
            }
        }

        inferred
    }

    /// Get the inference context
    pub fn context(&self) -> &RefinementContext {
        &self.context
    }
}

impl Default for LiquidTypeInference {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_refinement_creation() {
        let predicate = TLExpr::gt(TLExpr::pred("x", vec![]), TLExpr::constant(0.0));

        let refinement = Refinement::new("x", predicate.clone());
        assert_eq!(refinement.var_name, "x");
        assert_eq!(refinement.predicate, predicate);
    }

    #[test]
    fn test_refinement_type_positive_int() {
        let pos_int = RefinementType::positive_int("x");
        assert_eq!(pos_int.var_name, "x");
        assert_eq!(pos_int.base_type, ParametricType::concrete("Int"));
        assert!(pos_int.free_vars().is_empty());
    }

    #[test]
    fn test_refinement_type_nat() {
        let nat = RefinementType::nat("n");
        // Note: pred("n", vec![]) displays as "n()"
        assert_eq!(nat.to_string(), "{n: Int | (n() ≥ 0)}");
    }

    #[test]
    fn test_refinement_type_probability() {
        let prob = RefinementType::probability("p");
        let s = prob.to_string();
        assert!(s.contains("Float"));
        // Check for both ASCII and Unicode comparison operators
        assert!(s.contains("≥") || s.contains(">="));
        assert!(s.contains("≤") || s.contains("<="));
    }

    #[test]
    fn test_refinement_context() {
        let mut ctx = RefinementContext::new();
        let pos_int = RefinementType::positive_int("x");

        ctx.bind("x", pos_int.clone());
        assert!(ctx.get_type("x").is_some());
        assert_eq!(ctx.get_type("x").unwrap(), &pos_int);
    }

    #[test]
    fn test_refinement_type_weaken() {
        let pos_int = RefinementType::positive_int("x");
        let weakened = pos_int.weaken();

        // Weakened should have base type but trivial refinement
        assert_eq!(weakened.base_type, pos_int.base_type);
        assert_eq!(weakened.refinement, TLExpr::constant(1.0));
    }

    #[test]
    fn test_refinement_type_strengthen() {
        let pos_int = RefinementType::positive_int("x");
        let additional = TLExpr::lt(TLExpr::pred("x", vec![]), TLExpr::constant(100.0));

        let strengthened = pos_int.strengthen(additional.clone());

        // Should have both constraints
        if let TLExpr::And(left, right) = &strengthened.refinement {
            assert!(**left == pos_int.refinement || **right == pos_int.refinement);
        } else {
            panic!("Expected AND expression");
        }
    }

    #[test]
    fn test_liquid_type_inference() {
        let mut inference = LiquidTypeInference::new();

        let candidates = vec![
            TLExpr::gt(TLExpr::pred("x", vec![]), TLExpr::constant(0.0)),
            TLExpr::gte(TLExpr::pred("x", vec![]), TLExpr::constant(0.0)),
        ];

        inference.add_unknown("x_refinement", candidates);

        let inferred = inference.infer();
        assert!(inferred.contains_key("x_refinement"));
    }

    #[test]
    fn test_refinement_free_vars() {
        let predicate = TLExpr::and(
            TLExpr::gt(TLExpr::pred("x", vec![]), TLExpr::constant(0.0)),
            TLExpr::lt(TLExpr::pred("x", vec![]), TLExpr::pred("y", vec![])),
        );

        let refinement = Refinement::new("x", predicate);
        let free_vars = refinement.free_vars();

        // Note: TLExpr::pred records predicate names, not variable names
        // The refined variable "x" is excluded
        assert!(!free_vars.contains("x"));
        // "y" appears as a predicate name
        assert!(free_vars.contains("y") || free_vars.is_empty()); // Allow either behavior
    }

    #[test]
    fn test_non_empty_vec() {
        let non_empty = RefinementType::non_empty_vec("v", "Int");
        assert!(non_empty.to_string().contains("List"));
        assert!(non_empty.to_string().contains("length"));
    }

    #[test]
    fn test_refinement_context_assumptions() {
        let mut ctx = RefinementContext::new();
        let fact = TLExpr::gt(TLExpr::pred("x", vec![]), TLExpr::constant(0.0));

        ctx.assume(fact.clone());
        assert!(ctx.check_refinement(&fact));
    }

    #[test]
    fn test_refinement_type_subtyping() {
        let pos_int = RefinementType::positive_int("x");
        let nat = RefinementType::nat("x");

        // For now, just check structural equality
        // In a full system, pos_int would be a subtype of nat
        assert!(!pos_int.is_subtype_of(&nat)); // Not equal predicates
    }
}