haloumi-picus 0.5.6

use crate::pcl::display::{TextRepresentable, TextRepresentation};
use crate::pcl::errors::ExprArgsError;
use crate::pcl::expr::traits::ConstraintExpr;
use crate::pcl::vars::VarStr;
use anyhow::{Result, anyhow};
use haloumi_ir::Prime;
use impls::{AssumeDeterministicStmt, CallStmt, CommentLine, ConstraintStmt, PostConditionStmt};
use std::collections::HashSet;
use std::{cell::RefCell, rc::Rc};
use traits::{
    CallLikeAdaptor, ConstraintLike, ExprArgs, FreeVars, MaybeCallLike, StmtConstantFolding,
    StmtLike,
};

use crate::pcl::expr::Expr;

mod impls;
pub mod traits;

type Wrap<T> = Rc<RefCell<T>>;

pub type Stmt = Wrap<dyn StmtLike>;

impl<S: ExprArgs + ?Sized> ExprArgs for Wrap<S> {
    fn args(&self) -> Vec<Expr> {
        unsafe { (*self.as_ptr()).args() }
    }

    fn replace_arg(&mut self, idx: usize, expr: Expr) -> Result<(), ExprArgsError> {
        self.borrow_mut().replace_arg(idx, expr)
    }
}

impl<S: ConstraintLike + ?Sized> ConstraintLike for Wrap<S> {
    fn is_constraint(&self) -> bool {
        self.borrow().is_constraint()
    }

    fn constraint_expr(&self) -> Option<&dyn ConstraintExpr> {
        unsafe { (*self.as_ptr()).constraint_expr() }
    }
}

impl<S: MaybeCallLike + ?Sized> MaybeCallLike for Wrap<S> {
    fn as_call<'a>(&'a self) -> Option<CallLikeAdaptor<'a>> {
        unsafe { (*self.as_ptr()).as_call() }
    }
}

impl<S: StmtConstantFolding + ?Sized> StmtConstantFolding for Wrap<S> {
    fn fold(&self, prime: Prime) -> Option<Stmt> {
        self.borrow().fold(prime)
    }
}

impl<S: TextRepresentable + ?Sized> TextRepresentable for Wrap<S> {
    fn to_repr(&self) -> TextRepresentation<'_> {
        unsafe { (*self.as_ptr()).to_repr() }
    }

    fn width_hint(&self) -> usize {
        self.borrow().width_hint()
    }
}

impl<S: FreeVars + ?Sized> FreeVars for Wrap<S> {
    fn free_vars(&self) -> HashSet<&VarStr> {
        unsafe { (*self.as_ptr()).free_vars() }
    }
}

impl<T> StmtLike for Wrap<T> where T: StmtLike + PartialEq + ?Sized {}

//===----------------------------------------------------------------------===//
// Factories
//===----------------------------------------------------------------------===//

pub fn call(callee: String, inputs: Vec<Expr>, outputs: Vec<Expr>) -> Result<Stmt> {
    Ok(Wrap::new(
        CallStmt::new(
            callee,
            inputs,
            outputs
                .into_iter()
                .map(|e| {
                    e.var_name()
                        .cloned()
                        .ok_or_else(|| anyhow!("Output expressions can only be variable names"))
                })
                .collect::<Result<Vec<_>>>()?,
        )
        .into(),
    ))
}

pub fn assume_deterministic(expr: Expr) -> Result<Stmt> {
    // Manually unwrapped and wrapped for coercing to the right type
    Ok(expr
        .var_name()
        .map(AssumeDeterministicStmt::new)
        .map(Into::into)
        .map(Wrap::new)
        .ok_or_else(|| anyhow!("assume-deterministic argument must be a variable name"))?)
}

pub fn constrain(expr: Expr) -> Stmt {
    Wrap::new(ConstraintStmt::new(expr).into())
}

pub fn comment(s: String) -> Stmt {
    Wrap::new(CommentLine::new(s).into())
}

/// Creates a `(post-condition <expr>)` statement.
pub fn post_condition(expr: Expr) -> Stmt {
    Wrap::new(PostConditionStmt::new(expr).into())
}