nixjit/src/ty/internal/mod.rs

use std::cell::Cell;
use std::hash::Hash;
use std::marker::PhantomData;
use std::mem::MaybeUninit;
use std::ops::Deref;

use derive_more::{IsVariant, Unwrap};
use gc_arena::Mutation;
use gc_arena::lock::RefLock;
use gc_arena::{Collect, Gc, lock::GcRefLock};
use hashbrown::HashSet;

use super::common::*;
use super::public as p;

use crate::bytecode::OpCodes;
use crate::env::VmEnv;
use crate::error::*;
use crate::vm::VM;

mod attrset;
mod func;
mod list;
mod primop;
mod string;

pub use attrset::*;
pub use func::*;
pub use list::List;
pub use primop::*;

#[derive(Collect)]
#[collect(unsafe_drop)]
pub struct CoW<'gc, T: Clone + Collect<'gc>> {
    inner: Gc<'gc, CoWInner<'gc, T>>,
}

struct CoWInner<'gc, T: Clone + Collect<'gc>> {
    ref_count: Cell<usize>,
    data: MaybeUninit<T>,
    _marker: PhantomData<Cell<&'gc ()>>,
}

unsafe impl<'gc, T: Clone + Collect<'gc>> Collect<'gc> for CoWInner<'gc, T> {
    fn trace<Tr: gc_arena::collect::Trace<'gc>>(&self, cc: &mut Tr) {
        unsafe { self.data.assume_init_ref() }.trace(cc);
    }
}

#[allow(mutable_transmutes)]
impl<'gc, T: Clone + Collect<'gc>> CoW<'gc, T> {
    pub fn new(data: T, mc: &Mutation<'gc>) -> Self {
        Self {
            inner: Gc::new(mc, CoWInner::new(data)),
        }
    }

    pub fn new_cyclic(datafn: impl FnOnce(Self) -> T, mc: &Mutation<'gc>) -> Self {
        let inner = Gc::new(
            mc,
            CoWInner {
                ref_count: Cell::new(1),
                data: MaybeUninit::uninit(),
                _marker: PhantomData,
            },
        );
        let data = datafn(CoW { inner });
        unsafe { std::mem::transmute::<_, &mut MaybeUninit<_>>(&inner.data) }.write(data);
        Self { inner }
    }

    pub fn make_mut(&mut self, mc: &Mutation<'gc>) -> &mut T {
        if self.inner.ref_count.get() == 1 {
            unsafe { std::mem::transmute(self.inner.data.assume_init_ref()) }
        } else {
            unsafe {
                *self = CoW::new(self.inner.data.assume_init_ref().clone(), mc);
                std::mem::transmute(self.inner.data.assume_init_ref())
            }
        }
    }

    pub fn as_ptr(&self) -> *const T {
        self.as_ref() as *const T
    }
}

impl<'gc, T: Clone + Collect<'gc>> AsRef<T> for CoW<'gc, T> {
    fn as_ref(&self) -> &T {
        unsafe { self.inner.data.assume_init_ref() }
    }
}

impl<'gc, T: Clone + Collect<'gc>> Deref for CoW<'gc, T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        self.as_ref()
    }
}

impl<'gc, T: Clone + Collect<'gc>> Clone for CoW<'gc, T> {
    fn clone(&self) -> Self {
        self.inner.ref_count.set(self.inner.ref_count.get() + 1);
        Self { inner: self.inner }
    }
}

impl<'gc, T: Clone + Collect<'gc>> Drop for CoW<'gc, T> {
    fn drop(&mut self) {
        self.inner.ref_count.set(self.inner.ref_count.get() - 1);
    }
}

impl<'gc, T: Clone + Collect<'gc>> CoWInner<'gc, T> {
    fn new(data: T) -> Self {
        Self {
            ref_count: Cell::new(1),
            data: MaybeUninit::new(data),
            _marker: PhantomData,
        }
    }
}

#[derive(IsVariant, Unwrap, Clone, Collect)]
#[collect(no_drop)]
pub enum Value<'gc> {
    Int(i64),
    Float(f64),
    Bool(bool),
    String(CoW<'gc, String>),
    Null,
    Thunk(Thunk<'gc>),
    AttrSet(CoW<'gc, AttrSet<'gc>>),
    List(CoW<'gc, List<'gc>>),
    Catchable(Gc<'gc, String>),
    PrimOp(Gc<'gc, PrimOp>),
    PartialPrimOp(CoW<'gc, PartialPrimOp<'gc>>),
    Func(Gc<'gc, Func<'gc>>),
}

impl Hash for Value<'_> {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        use Value::*;
        std::mem::discriminant(self).hash(state);
        match self {
            Int(x) => x.hash(state),
            Float(x) => x.to_bits().hash(state),
            Bool(x) => x.hash(state),
            Null => (),
            String(x) => x.as_ptr().hash(state),
            Thunk(x) => (x as *const self::Thunk).hash(state),
            AttrSet(x) => x.as_ptr().hash(state),
            List(x) => x.as_ptr().hash(state),
            Catchable(x) => x.hash(state),
            PrimOp(x) => (x.as_ref() as *const self::PrimOp).hash(state),
            PartialPrimOp(x) => (x.as_ref() as *const self::PartialPrimOp).hash(state),
            Func(x) => (x.as_ref() as *const self::Func).hash(state),
        }
    }
}

impl<'gc> Value<'gc> {
    fn eq_impl(&self, other: &Self) -> bool {
        use Value::*;
        match (self, other) {
            (Bool(a), Bool(b)) => a == b,
            (Int(a), Int(b)) => a == b,
            (Float(a), Float(b)) => a == b,
            (Int(a), Float(b)) => *a as f64 == *b,
            (Float(a), Int(b)) => *b as f64 == *a,
            (String(a), String(b)) => a.as_str().eq(b.as_str()),
            (Null, Null) => true,
            (AttrSet(a), AttrSet(b)) => a.eq_impl(b),
            (List(a), List(b)) => a.eq(b),
            _ => false,
        }
    }
}

impl<'jit: 'vm, 'vm, 'gc> PartialEq for Value<'gc> {
    fn eq(&self, other: &Self) -> bool {
        use Value::*;
        match (self, other) {
            (AttrSet(a), AttrSet(b)) => a.as_ptr().eq(&b.as_ptr()),
            (List(a), List(b)) => a.as_ptr().eq(&b.as_ptr()),
            _ => false,
        }
    }
}

impl Eq for Value<'_> {}

#[derive(IsVariant, Unwrap, Clone)]
pub enum ValueAsRef<'v, 'gc> {
    Int(i64),
    Float(f64),
    Bool(bool),
    String(&'v str),
    Null,
    Thunk(&'v Thunk<'gc>),
    AttrSet(&'v AttrSet<'gc>),
    List(&'v List<'gc>),
    Catchable(&'v str),
    PrimOp(&'v PrimOp),
    PartialPrimOp(&'v PartialPrimOp<'gc>),
    Func(&'v Func<'gc>),
}

#[derive(IsVariant, Unwrap)]
pub enum ValueAsMut<'v, 'gc> {
    Int(i64),
    Float(f64),
    Bool(bool),
    String(&'v str),
    Null,
    Thunk(&'v Thunk<'gc>),
    AttrSet(&'v mut AttrSet<'gc>),
    List(&'v mut List<'gc>),
    Catchable(&'v str),
    PrimOp(&'v PrimOp),
    PartialPrimOp(&'v mut PartialPrimOp<'gc>),
    Func(&'v Func<'gc>),
}

impl<'gc, 'v> Value<'gc> {
    pub fn as_ref(&'v self) -> ValueAsRef<'v, 'gc> {
        use Value::*;
        use ValueAsRef as R;
        match self {
            Int(x) => R::Int(*x),
            Float(x) => R::Float(*x),
            Bool(x) => R::Bool(*x),
            String(x) => R::String(x),
            Null => R::Null,
            Thunk(x) => R::Thunk(x),
            AttrSet(x) => R::AttrSet(x),
            List(x) => R::List(x),
            Catchable(x) => R::Catchable(x),
            PrimOp(x) => R::PrimOp(x),
            PartialPrimOp(x) => R::PartialPrimOp(x),
            Func(x) => R::Func(x),
        }
    }

    pub fn as_mut(&'v mut self, mc: &Mutation<'gc>) -> ValueAsMut<'v, 'gc> {
        use Value::*;
        use ValueAsMut as M;
        match self {
            Int(x) => M::Int(*x),
            Float(x) => M::Float(*x),
            Bool(x) => M::Bool(*x),
            String(x) => M::String(x),
            Null => M::Null,
            Thunk(x) => M::Thunk(x),
            AttrSet(x) => M::AttrSet(x.make_mut(mc)),
            List(x) => M::List(x.make_mut(mc)),
            Catchable(x) => M::Catchable(x),
            PrimOp(x) => M::PrimOp(x),
            PartialPrimOp(x) => M::PartialPrimOp(x.make_mut(mc)),
            Func(x) => M::Func(x),
        }
    }
}
impl<'gc> Value<'gc> {
    pub fn ok(self) -> Result<Self> {
        Ok(self)
    }

    pub fn typename(&self) -> &'static str {
        use Value::*;
        match self {
            Int(_) => "int",
            Float(_) => "float",
            Bool(_) => "bool",
            String(_) => "string",
            Null => "null",
            Thunk(_) => "thunk",
            AttrSet(_) => "set",
            List(_) => "list",
            Catchable(_) => unreachable!(),
            PrimOp(_) => "lambda",
            PartialPrimOp(_) => "lambda",
            Func(_) => "lambda",
        }
    }

    pub fn callable(&self) -> bool {
        match self {
            Value::PrimOp(_) | Value::PartialPrimOp(_) | Value::Func(_) => true,
            Value::AttrSet(_) => todo!(),
            _ => false,
        }
    }

    pub fn call(&mut self, arg: Self, vm: &VM, mc: &Mutation<'gc>) -> Result<()> {
        use Value::*;
        if matches!(arg, Value::Catchable(_)) {
            *self = arg;
            return Ok(());
        }
        *self = match self {
            PrimOp(func) => func.call(arg, vm, mc),
            PartialPrimOp(func) => func.call(arg, vm, mc),
            // Value::Func(func) => func.call(arg, vm, mc),
            Catchable(_) => return Ok(()),
            _ => todo!(),
        }?;
        Ok(())
    }

    pub fn not(&mut self) {
        use Value::*;
        *self = match &*self {
            Bool(bool) => Bool(!bool),
            Value::Catchable(_) => return,
            _ => todo!(),
        }
    }

    pub fn and(&mut self, other: Self) {
        use Value::*;
        *self = match (&*self, other) {
            (Bool(a), Bool(b)) => Bool(*a && b),
            (Value::Catchable(_), _) => return,
            (_, x @ Value::Catchable(_)) => x,
            _ => todo!(),
        }
    }

    pub fn or(&mut self, other: Self) {
        use Value::*;
        *self = match (&*self, other) {
            (Bool(a), Bool(b)) => Bool(*a || b),
            (Value::Catchable(_), _) => return,
            (_, x @ Value::Catchable(_)) => x,
            _ => todo!(),
        }
    }

    pub fn eq(&mut self, other: Self) {
        use Value::Bool;
        *self = match (&*self, other) {
            (Value::Catchable(_), _) => return,
            (_, x @ Value::Catchable(_)) => x,
            (s, other) => Bool(s.eq_impl(&other)),
        };
    }

    pub fn lt(&mut self, other: Self) {
        use Value::*;
        *self = Bool(match (&*self, other) {
            (Int(a), Int(b)) => *a < b,
            (Int(a), Float(b)) => (*a as f64) < b,
            (Float(a), Int(b)) => *a < b as f64,
            (Float(a), Float(b)) => *a < b,
            (String(a), String(b)) => a.as_str() < b.as_str(),
            (Value::Catchable(_), _) => return,
            (_, x @ Value::Catchable(_)) => {
                *self = x;
                return;
            }
            _ => todo!(),
        })
    }

    pub fn neg(&mut self) {
        use Value::*;
        *self = match &*self {
            Int(int) => Int(-int),
            Float(float) => Float(-float),
            Value::Catchable(_) => return,
            _ => todo!(),
        }
    }

    pub fn add(&mut self, other: Self, mc: &Mutation<'gc>) {
        use Value::*;
        *self = match (&*self, other) {
            (Int(a), Int(b)) => Int(a + b),
            (Int(a), Float(b)) => Float(*a as f64 + b),
            (Float(a), Int(b)) => Float(a + b as f64),
            (Float(a), Float(b)) => Float(a + b),
            (String(a), String(b)) => {
                let mut a = a.clone();
                a.make_mut(mc).push_str(b.as_str());
                String(a)
            }
            (Value::Catchable(_), _) => return,
            (_, x @ Value::Catchable(_)) => x,
            _ => todo!(),
        }
    }

    pub fn mul(&mut self, other: Self) {
        use Value::*;
        *self = match (&*self, other) {
            (Int(a), Int(b)) => Int(a * b),
            (Int(a), Float(b)) => Float(*a as f64 * b),
            (Float(a), Int(b)) => Float(a * b as f64),
            (Float(a), Float(b)) => Float(a * b),
            (Value::Catchable(_), _) => return,
            (_, x @ Value::Catchable(_)) => x,
            _ => todo!(),
        }
    }

    pub fn div(&mut self, other: Self) -> Result<()> {
        use Value::*;
        *self = match (&*self, other) {
            (_, Int(0)) => return Err(Error::EvalError("division by zero".to_string())),
            (_, Float(0.)) => {
                return Err(Error::EvalError("division by zero".to_string()));
            }
            (Int(a), Int(b)) => Int(a / b),
            (Int(a), Float(b)) => Float(*a as f64 / b),
            (Float(a), Int(b)) => Float(a / b as f64),
            (Float(a), Float(b)) => Float(a / b),
            (Value::Catchable(_), _) => return Ok(()),
            (_, x @ Value::Catchable(_)) => x,
            _ => todo!(),
        };
        Ok(())
    }

    pub fn concat_string(&mut self, mut other: Self, mc: &Mutation<'gc>) -> &mut Self {
        match (self.coerce_to_string(), other.coerce_to_string()) {
            (Value::String(a), Value::String(b)) => {
                let mut a = a.clone();
                a.make_mut(mc).push_str(b.as_str())
            },
            (_, Value::Catchable(_)) => *self = other,
            (Value::Catchable(_), _) => (),
            _ => todo!(),
        }
        self
    }

    pub fn push(&mut self, elem: Self, mc: &Mutation<'gc>) -> &mut Self {
        if let Value::List(list) = self {
            list.make_mut(mc).push(elem);
        } else if let Value::Catchable(_) = self {
        } else if let Value::Catchable(_) = elem {
            *self = elem;
        } else {
            todo!()
        }
        self
    }

    pub fn concat(&mut self, other: Self, mc: &Mutation<'gc>) {
        if let x @ Value::Catchable(_) = other {
            *self = x;
            return;
        }
        match (self, other) {
            (Value::List(a), Value::List(b)) => {
                a.make_mut(mc).concat(b.as_ref());
            }
            (Value::Catchable(_), _) => (),
            _ => todo!(),
        }
    }

    pub fn push_attr(&mut self, sym: usize, val: Self, mc: &Mutation<'gc>) -> &mut Self {
        if let Value::AttrSet(attrs) = self {
            attrs.make_mut(mc).push_attr(sym, val);
        } else if let Value::Catchable(_) = self {
        } else if let Value::Catchable(_) = val {
            *self = val
        } else {
            todo!()
        }
        self
    }

    pub fn update(&mut self, other: Self, mc: &Mutation<'gc>) {
        if let x @ Value::Catchable(_) = other {
            *self = x;
            return;
        }
        match (self, other) {
            (Value::AttrSet(a), Value::AttrSet(b)) => {
                a.make_mut(mc).update(b.as_ref());
            }
            (Value::Catchable(_), _) => (),
            _ => todo!(),
        }
    }

    pub fn select(&mut self, sym: usize, vm: &VM<'gc>) -> Result<&mut Self> {
        let val = match self {
            Value::AttrSet(attrs) => attrs
                .select(sym)
                .ok_or_else(|| Error::EvalError(format!("{} not found", vm.get_sym(sym)))),
            Value::Catchable(_) => return Ok(self),
            _ => Err(Error::EvalError(format!(
                "cannot select from {:?}",
                self.typename()
            ))),
        }?;
        *self = val;
        Ok(self)
    }

    pub fn select_with_default(&mut self, sym: usize, default: Self) -> Result<&mut Self> {
        let val = match self {
            Value::AttrSet(attrs) => attrs.select(sym).unwrap_or(default),
            Value::Catchable(_) => return Ok(self),
            _ => {
                return Err(Error::EvalError(format!(
                    "cannot select from {:?}",
                    self.typename()
                )));
            }
        };
        *self = val;
        Ok(self)
    }

    pub fn has_attr(&mut self, sym: usize) -> &mut Self {
        if let Value::AttrSet(attrs) = self {
            let val = Value::Bool(attrs.has_attr(sym));
            *self = val;
        } else if let Value::Catchable(_) = self {
        } else {
            *self = Value::Bool(false);
        }
        self
    }

    pub fn coerce_to_string(&mut self) -> &mut Self {
        if let Value::String(_) = self {
        } else if let Value::Catchable(_) = self {
        } else {
            todo!()
        }
        self
    }

    pub fn to_public(&self, vm: &VM<'gc>, seen: &mut HashSet<Value<'gc>>) -> p::Value {
        use self::Value::*;
        use p::Value;
        if seen.contains(self) {
            return Value::Repeated;
        }
        match self {
            AttrSet(attrs) => {
                seen.insert(self.clone());
                attrs.to_public(vm, seen)
            }
            List(list) => {
                seen.insert(self.clone());
                list.to_public(vm, seen)
            }
            Catchable(catchable) => Value::Catchable(catchable.as_ref().clone().into()),
            Int(x) => Value::Const(Const::Int(*x)),
            Float(x) => Value::Const(Const::Float(*x)),
            Bool(x) => Value::Const(Const::Bool(*x)),
            String(x) => Value::Const(Const::String(x.as_ref().into())),
            Null => Value::Const(Const::Null),
            Thunk(_) => Value::Thunk,
            PrimOp(primop) => Value::PrimOp(primop.name),
            PartialPrimOp(primop) => Value::PartialPrimOp(primop.name),
            Func(_) => Value::Func,
        }
    }
}

#[repr(transparent)]
#[derive(Clone, Collect)]
#[collect(no_drop)]
pub struct Thunk<'gc> {
    pub thunk: GcRefLock<'gc, _Thunk<'gc>>,
}

#[derive(IsVariant, Unwrap)]
pub enum _Thunk<'gc> {
    Code(&'gc OpCodes, Option<Gc<'gc, VmEnv<'gc>>>),
    Suspended,
    Value(Value<'gc>),
}

unsafe impl<'gc> Collect<'gc> for _Thunk<'gc> {
    fn trace<T: gc_arena::collect::Trace<'gc>>(&self, cc: &mut T) {
        use _Thunk::*;
        match self {
            Code(_, env) => env.trace(cc),
            Value(val) => val.trace(cc),
            _ => (),
        }
    }
}

impl<'gc> Thunk<'gc> {
    pub fn new(opcodes: &'gc OpCodes, mc: &Mutation<'gc>) -> Self {
        Thunk {
            thunk: Gc::new(mc, RefLock::new(_Thunk::Code(opcodes, None))),
        }
    }

    pub fn capture_env(&self, env: Gc<'gc, VmEnv<'gc>>, mc: &Mutation<'gc>) {
        if let _Thunk::Code(_, envcell) = &mut *self.thunk.borrow_mut(mc) {
            let _ = envcell.insert(env);
        }
    }

    pub fn suspend(&self, mc: &Mutation<'gc>) -> Result<(&'gc OpCodes, Gc<'gc, VmEnv<'gc>>)> {
        let _Thunk::Code(opcodes, env) =
            std::mem::replace(&mut *self.thunk.borrow_mut(mc), _Thunk::Suspended)
        else {
            return Err(Error::EvalError("infinite recursion occured".into()))
        };
        Ok((opcodes, env.unwrap()))
    }

    pub fn insert_value(&self, value: Value<'gc>, mc: &Mutation<'gc>) {
        *self.thunk.borrow_mut(mc) = _Thunk::Value(value);
    }

    pub fn get_value(&self) -> Option<Value<'gc>> {
        if let _Thunk::Value(val) = &*self.thunk.borrow() {
            Some(val.clone())
        } else {
            None
        }
    }
}

impl PartialEq for Thunk<'_> {
    fn eq(&self, _: &Self) -> bool {
        false
    }
}