xlsynth 0.44.0

// SPDX-License-Identifier: Apache-2.0

use std::fmt;
use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};

use crate::ir_analysis::{IrAnalysis, IrAnalysisLevel};
use crate::lib_support::{self, c_str_to_rust, xls_function_jit_run, xls_make_function_jit};
pub use crate::lib_support::{RunResult, TraceMessage};
use crate::xlsynth_error::XlsynthError;
use crate::{
    lib_support::{
        xls_function_get_name, xls_function_get_type, xls_function_to_string,
        xls_function_type_to_string, xls_interpret_function, xls_package_free,
        xls_package_get_function, xls_package_get_functions, xls_package_get_type_for_value,
        xls_package_to_string, xls_parse_ir_package, xls_type_to_string,
    },
    IrValue,
};
use xlsynth_sys::{CIrFunction, CIrPackage, CScheduleAndCodegenResult};

// -- ScheduleAndCodegenResult

pub struct ScheduleAndCodegenResult {
    pub(crate) ptr: *mut CScheduleAndCodegenResult,
}

impl Drop for ScheduleAndCodegenResult {
    fn drop(&mut self) {
        if !self.ptr.is_null() {
            unsafe {
                xlsynth_sys::xls_schedule_and_codegen_result_free(self.ptr);
            }
            self.ptr = std::ptr::null_mut();
        }
    }
}

impl ScheduleAndCodegenResult {
    pub fn get_verilog_text(&self) -> Result<String, XlsynthError> {
        unsafe {
            let verilog = xlsynth_sys::xls_schedule_and_codegen_result_get_verilog_text(self.ptr);
            let verilog_str = c_str_to_rust(verilog);
            Ok(verilog_str)
        }
    }
}

// -- IrPackage & IrPackagePtr

/// We wrap up the raw C pointer in this type that implements `Drop` trait --
/// this allows us to wrap it up in an Arc so that types with derived lifetimes
/// (e.g. C function pointers) can grab a hold of the Arc to prevent
/// deallocation of its backing type.
pub(crate) struct IrPackagePtr(pub *mut CIrPackage);

impl IrPackagePtr {
    pub fn mut_c_ptr(&self) -> *mut CIrPackage {
        self.0
    }
    pub fn const_c_ptr(&self) -> *const CIrPackage {
        self.0
    }
}

impl Drop for IrPackagePtr {
    fn drop(&mut self) {
        if !self.0.is_null() {
            xls_package_free(self.0);
            self.0 = std::ptr::null_mut();
        }
    }
}

pub struct IrPackage {
    pub(crate) ptr: Arc<RwLock<IrPackagePtr>>,
    // TODO(cdleary): 2024-06-23 This is the filename that is passed to the "parse package IR"
    // C API functionality. We should be able to remove this field and recover the value
    // from the built package if we had the appropriate C API.
    pub(crate) filename: Option<String>,
}

unsafe impl Send for IrPackage {}
unsafe impl Sync for IrPackage {}

impl IrPackage {
    fn with_read<T>(&self, f: impl FnOnce(RwLockReadGuard<IrPackagePtr>) -> T) -> T {
        f(self.ptr.read().unwrap())
    }

    fn with_write<T>(&self, f: impl FnOnce(RwLockWriteGuard<IrPackagePtr>) -> T) -> T {
        f(self.ptr.write().unwrap())
    }

    pub fn new(name: &str) -> Result<Self, XlsynthError> {
        lib_support::xls_package_new(name)
    }

    pub fn parse_ir(ir: &str, filename: Option<&str>) -> Result<Self, XlsynthError> {
        xls_parse_ir_package(ir, filename)
    }

    pub fn parse_ir_from_path(path: &std::path::Path) -> Result<Self, XlsynthError> {
        let ir = match std::fs::read_to_string(path) {
            Ok(ir) => ir,
            Err(e) => {
                return Err(XlsynthError(format!("Failed to read IR from path: {e}")));
            }
        };
        let filename = path.file_name().and_then(|s| s.to_str());
        let ir_package = Self::parse_ir(&ir, filename)?;
        Ok(ir_package)
    }

    pub fn verify(&self) -> Result<(), XlsynthError> {
        self.with_write(|guard| lib_support::xls_verify_package(guard.mut_c_ptr()))
    }

    pub fn set_top_by_name(&mut self, name: &str) -> Result<(), XlsynthError> {
        self.with_write(|guard| lib_support::xls_package_set_top_by_name(guard.mut_c_ptr(), name))
    }

    /// Creates an IR analysis handle for this package.
    ///
    /// Precondition: this package must have a top function base set (e.g. via
    /// [`IrPackage::set_top_by_name`]).
    pub fn create_ir_analysis(&self) -> Result<IrAnalysis, XlsynthError> {
        IrAnalysis::create_from_package_ptr(self.ptr.clone())
    }

    /// Creates an IR analysis handle for this package with the given analysis
    /// level.
    ///
    /// Precondition: this package must have a top function base set (e.g. via
    /// [`IrPackage::set_top_by_name`]).
    pub fn create_ir_analysis_with_level(
        &self,
        level: IrAnalysisLevel,
    ) -> Result<IrAnalysis, XlsynthError> {
        IrAnalysis::create_from_package_ptr_with_level(self.ptr.clone(), level)
    }

    pub fn get_function(&self, name: &str) -> Result<IrFunction, XlsynthError> {
        self.with_read(|guard| xls_package_get_function(&self.ptr, guard, name))
    }

    /// Returns all functions contained within this package.
    ///
    /// The returned [`IrFunction`] values are tied to this package's lifetime;
    /// they do not need to be freed separately.
    pub fn get_functions(&self) -> Result<Vec<IrFunction>, XlsynthError> {
        self.with_write(|guard| xls_package_get_functions(&self.ptr, guard))
    }

    pub fn get_type_for_value(&self, value: &IrValue) -> Result<IrType, XlsynthError> {
        let parent = self.ptr.clone();
        self.with_write(|guard| {
            let ptr = xls_package_get_type_for_value(guard.mut_c_ptr(), value.ptr)?;
            Ok(IrType::new(ptr, parent))
        })
    }

    pub fn get_bits_type(&self, bit_count: u64) -> IrType {
        let parent = self.ptr.clone();
        self.with_write(|guard| {
            let ptr = lib_support::xls_package_get_bits_type(guard.mut_c_ptr(), bit_count);
            IrType::new(ptr, parent)
        })
    }

    pub fn get_tuple_type(&self, members: &[IrType]) -> IrType {
        let parent = self.ptr.clone();
        self.with_write(|guard| {
            let ptr = lib_support::xls_package_get_tuple_type(guard.mut_c_ptr(), members);
            IrType::new(ptr, parent)
        })
    }

    pub fn types_eq(&self, a: &IrType, b: &IrType) -> Result<bool, XlsynthError> {
        self.with_read(|_| Ok(a.ptr == b.ptr))
    }

    pub fn get_token_type(&self) -> IrType {
        let parent = self.ptr.clone();
        self.with_write(|guard| {
            let ptr = lib_support::xls_package_get_token_type(guard.mut_c_ptr());
            IrType::new(ptr, parent)
        })
    }

    pub fn get_array_type(&self, element_type: &IrType, size: i64) -> IrType {
        let parent = self.ptr.clone();
        self.with_write(|guard| {
            let ptr =
                lib_support::xls_package_get_array_type(guard.mut_c_ptr(), element_type.ptr, size);
            IrType::new(ptr, parent)
        })
    }

    pub fn filename(&self) -> Option<&str> {
        self.filename.as_deref()
    }
}

pub struct IrType {
    pub(crate) ptr: *mut xlsynth_sys::CIrType,
    pub(crate) _parent: Arc<RwLock<IrPackagePtr>>,
}

impl std::fmt::Display for IrType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", xls_type_to_string(self.ptr).unwrap())
    }
}

impl IrType {
    pub(crate) fn new(ptr: *mut xlsynth_sys::CIrType, parent: Arc<RwLock<IrPackagePtr>>) -> Self {
        Self {
            ptr,
            _parent: parent,
        }
    }

    pub fn get_flat_bit_count(&self) -> u64 {
        lib_support::xls_type_get_flat_bit_count(self.ptr)
    }
}

pub struct IrFunctionType {
    pub(crate) ptr: *mut xlsynth_sys::CIrFunctionType,
    pub(crate) parent: Arc<RwLock<IrPackagePtr>>,
}

impl std::fmt::Display for IrFunctionType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", xls_function_type_to_string(self.ptr).unwrap())
    }
}

impl IrFunctionType {
    pub(crate) fn new(
        ptr: *mut xlsynth_sys::CIrFunctionType,
        parent: Arc<RwLock<IrPackagePtr>>,
    ) -> Self {
        Self { ptr, parent }
    }

    /// Returns the number of parameters in this function type.
    pub fn param_count(&self) -> usize {
        lib_support::xls_function_type_param_count(self.ptr) as usize
    }

    /// Returns the type of the `i`th parameter.
    pub fn param_type(&self, index: usize) -> Result<IrType, XlsynthError> {
        let ptr = lib_support::xls_function_type_get_param_type(self.ptr, index)?;
        Ok(IrType::new(ptr, self.parent.clone()))
    }

    /// Returns the return type of the function.
    pub fn return_type(&self) -> IrType {
        let ptr = lib_support::xls_function_type_get_return_type(self.ptr);
        IrType::new(ptr, self.parent.clone())
    }
}

impl fmt::Display for IrPackage {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let s = self.with_read(|guard| xls_package_to_string(guard.const_c_ptr()).unwrap());
        write!(f, "{s}")
    }
}

pub struct IrFunction {
    pub(crate) parent: Arc<RwLock<IrPackagePtr>>,
    pub(crate) ptr: *mut CIrFunction,
}

unsafe impl Send for IrFunction {}
unsafe impl Sync for IrFunction {}

impl IrFunction {
    /// Creates an IR analysis handle for this function's owning package.
    ///
    /// Precondition: the owning package must have a top function base set (e.g.
    /// via [`IrPackage::set_top_by_name`]).
    pub fn create_ir_analysis(&self) -> Result<IrAnalysis, XlsynthError> {
        IrAnalysis::create_from_package_ptr(self.parent.clone())
    }

    /// Creates an IR analysis handle for this function's owning package with
    /// the given analysis level.
    ///
    /// Precondition: the owning package must have a top function base set (e.g.
    /// via [`IrPackage::set_top_by_name`]).
    pub fn create_ir_analysis_with_level(
        &self,
        level: IrAnalysisLevel,
    ) -> Result<IrAnalysis, XlsynthError> {
        IrAnalysis::create_from_package_ptr_with_level(self.parent.clone(), level)
    }

    pub fn interpret(&self, args: &[IrValue]) -> Result<IrValue, XlsynthError> {
        let package_read_guard: RwLockReadGuard<IrPackagePtr> = self.parent.read().unwrap();
        xls_interpret_function(&package_read_guard, self.ptr, args)
    }

    pub fn to_ir_string(&self) -> Result<String, XlsynthError> {
        // We take a read guard to document that `self.ptr` is owned by the
        // package and must not be used after the package is dropped.
        let _package_read_guard: RwLockReadGuard<IrPackagePtr> = self.parent.read().unwrap();
        xls_function_to_string(self.ptr)
    }

    /// Converts this function to Z3 SMT-LIB text.
    pub fn to_z3_smtlib(&self) -> Result<String, XlsynthError> {
        // We take a read guard to document that `self.ptr` is owned by the
        // package and must not be used after the package is dropped.
        let _package_read_guard: RwLockReadGuard<IrPackagePtr> = self.parent.read().unwrap();
        lib_support::xls_function_to_z3_smtlib(self.ptr)
    }

    pub fn get_name(&self) -> String {
        xls_function_get_name(self.ptr).unwrap()
    }

    pub fn get_type(&self) -> Result<IrFunctionType, XlsynthError> {
        let package_write_guard: RwLockWriteGuard<IrPackagePtr> = self.parent.write().unwrap();
        let ptr = xls_function_get_type(&package_write_guard, self.ptr)?;
        Ok(IrFunctionType::new(ptr, self.parent.clone()))
    }

    /// Returns the name of the `i`th parameter to this function.
    pub fn param_name(&self, index: usize) -> Result<String, XlsynthError> {
        lib_support::xls_function_get_param_name(self.ptr, index)
    }
}

impl fmt::Display for IrFunction {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.to_ir_string().unwrap())
    }
}

pub struct IrFunctionJit {
    pub(crate) parent: Arc<RwLock<IrPackagePtr>>,
    pub(crate) ptr: *mut xlsynth_sys::CIrFunctionJit,
}

impl Drop for IrFunctionJit {
    fn drop(&mut self) {
        if !self.ptr.is_null() {
            unsafe {
                xlsynth_sys::xls_function_jit_free(self.ptr);
            }
        }
    }
}

impl IrFunctionJit {
    pub fn new(function: &IrFunction) -> Result<Self, XlsynthError> {
        let package_read_guard: RwLockReadGuard<IrPackagePtr> = function.parent.read().unwrap();
        let ptr = xls_make_function_jit(&package_read_guard, function.ptr)?;
        Ok(IrFunctionJit {
            parent: function.parent.clone(),
            ptr,
        })
    }

    pub fn run(&self, args: &[IrValue]) -> Result<RunResult, XlsynthError> {
        let package_read_guard: RwLockReadGuard<IrPackagePtr> = self.parent.read().unwrap();
        xls_function_jit_run(&package_read_guard, self.ptr, args)
    }
}

#[cfg(test)]
mod tests {
    use crate::FnBuilder;

    use super::*;

    /// Verifies that formatting the parsed IR yields the exact original text.
    fn round_trip_ir_text(ir_text: &str) -> String {
        let package = IrPackage::parse_ir(ir_text, None).expect("parse success");
        let printed = package.to_string();
        assert_eq!(printed, ir_text);
        printed
    }

    #[test]
    fn test_ir_package_parse() {
        let ir =
            "package test\nfn f() -> bits[32] { ret literal.1: bits[32] = literal(value=42) }\n";
        let package = IrPackage::parse_ir(ir, None).expect("parse success");
        let f = package.get_function("f").expect("should find function");
        assert_eq!(f.get_name(), "f");
        let result = f.interpret(&[]).expect("interpret success");
        assert_eq!(result, IrValue::parse_typed("bits[32]:42").unwrap());
    }

    #[test]
    fn test_plus_one_fn_interp() {
        let ir = "package test\nfn plus_one(x: bits[32]) -> bits[32] {
    literal.2: bits[32] = literal(value=1)
    ret add.1: bits[32] = add(x, literal.2)
}";
        let package = IrPackage::parse_ir(ir, None).expect("parse success");
        let f = package
            .get_function("plus_one")
            .expect("should find function");
        assert_eq!(f.get_name(), "plus_one");

        // Inspect the function type.
        let f_type = f.get_type().expect("get type success");
        assert_eq!(f_type.to_string(), "(bits[32]) -> bits[32]".to_string());

        let ft = IrValue::parse_typed("bits[32]:42").unwrap();
        let result = f.interpret(&[ft]).expect("interpret success");
        let want = IrValue::parse_typed("bits[32]:43").unwrap();
        assert_eq!(result, want);

        assert_eq!(
            package.get_type_for_value(&want).unwrap().to_string(),
            "bits[32]".to_string()
        );
    }

    #[test]
    fn test_ir_function_to_string() {
        let ir = "package test\nfn plus_one(x: bits[32]) -> bits[32] {
    literal.2: bits[32] = literal(value=1)
    ret add.1: bits[32] = add(x, literal.2)
}\n";
        let package = IrPackage::parse_ir(ir, None).expect("parse success");
        let f = package
            .get_function("plus_one")
            .expect("should find function");

        let printed = f.to_ir_string().expect("stringify success");

        // Note: this is a golden string for the XLS IR printer for a single
        // function. It should be stable and deterministic.
        let expected = "fn plus_one(x: bits[32] id=4) -> bits[32] {\n  literal.2: bits[32] = literal(value=1, id=2)\n  ret add.1: bits[32] = add(x, literal.2, id=1)\n}\n";
        assert_eq!(printed, expected);
    }

    #[test]
    fn test_ir_package_set_top_by_name() {
        let mut package = IrPackage::new("test_package").unwrap();
        // Build the identity function inside of the package that is not marked as top.
        let u32 = package.get_bits_type(32);
        let mut builder = FnBuilder::new(&mut package, "f", true);
        let x = builder.param("x", &u32);
        let _f = builder.build_with_return_value(&x);

        assert_eq!(
            package.to_string(),
            "package test_package\n\nfn f(x: bits[32] id=1) -> bits[32] {
  ret x: bits[32] = param(name=x, id=1)
}
"
        );
        package.set_top_by_name("f").unwrap();
        assert_eq!(
            package.to_string(),
            "package test_package\n\ntop fn f(x: bits[32] id=1) -> bits[32] {
  ret x: bits[32] = param(name=x, id=1)
}
"
        );
    }

    #[test]
    fn test_ir_package_get_functions() {
        let mut package = IrPackage::new("test_package").unwrap();
        let u32 = package.get_bits_type(32);

        let mut builder = FnBuilder::new(&mut package, "alpha", true);
        let x = builder.param("x", &u32);
        builder.build_with_return_value(&x).unwrap();

        let mut builder = FnBuilder::new(&mut package, "beta", true);
        let y = builder.param("y", &u32);
        builder.build_with_return_value(&y).unwrap();

        let mut names: Vec<String> = package
            .get_functions()
            .unwrap()
            .into_iter()
            .map(|f| f.get_name())
            .collect();
        names.sort();
        assert_eq!(names, vec!["alpha".to_string(), "beta".to_string()]);
    }

    #[test]
    fn test_ir_types_keep_package_alive() {
        let ir = r#"package test

fn f(x: bits[8] id=1) -> bits[8] {
  k: bits[8] = literal(value=240, id=2)
  ret r: bits[8] = and(x, k, id=3)
}
"#;
        let mut package = IrPackage::parse_ir(ir, None).expect("parse success");
        package.set_top_by_name("f").expect("set top");
        let f = package.get_function("f").expect("get function");
        let f_type = f.get_type().expect("get function type");
        let token_type = package.get_token_type();
        drop(package);
        assert_eq!(f_type.to_string(), "(bits[8]) -> bits[8]");
        assert_eq!(token_type.to_string(), "token");
    }

    #[test]
    fn test_ir_type_tuple() {
        let package = IrPackage::new("test_package").unwrap();
        let u32 = package.get_bits_type(32);
        let u64 = package.get_bits_type(64);
        let tuple = package.get_tuple_type(&[u32, u64]);
        assert_eq!(tuple.to_string(), "(bits[32], bits[64])");
    }

    #[test]
    fn test_ir_package_verify_succeeds() {
        let ir = "package test\nfn f() -> bits[32] {\n  ret literal.1: bits[32] = literal(value=42)\n}\n";
        let package = IrPackage::parse_ir(ir, None).expect("parse success");
        package.verify().expect("verify success");
    }

    #[test]
    fn test_ir_package_verify_fails_when_builder_skips_checks() {
        let mut package = IrPackage::new("test_package").unwrap();
        let u32 = package.get_bits_type(32);

        // Construct a function with mismatched operand types using a builder that does
        // not verify nodes as they are added. The package verifier should catch
        // this mismatch.
        let mut builder = FnBuilder::new(&mut package, "bad", false);
        let x = builder.param("x", &u32);
        let wide_literal = IrValue::parse_typed("bits[64]:1").unwrap();
        let literal = builder.literal(&wide_literal, None);
        let sum = builder.add(&x, &literal, None);

        builder
            .build_with_return_value(&sum)
            .expect("builder should succeed without verification");

        let error = package.verify().expect_err("verify should fail");
        assert!(
            error.to_string().contains("type") || error.to_string().contains("Type"),
            "unexpected error: {}",
            error
        );
    }

    #[test]
    fn test_ir_round_trip_trace_and_cover_ops() {
        let ir = r#"package test

fn f(x: bits[1] id=7) -> bits[1] {
  after_all.2: token = after_all(id=2)
  literal.1: bits[1] = literal(value=1, id=1)
  trace.3: token = trace(after_all.2, x, format="x={}", data_operands=[x], id=3)
  cover.4: () = cover(x, label="x_is_one", id=4)
  ret literal.5: bits[1] = literal(value=1, id=5)
}
"#;

        round_trip_ir_text(ir);
    }
}