nixjit/src/ty/internal/mod.rs

use std::cell::RefCell;
use std::sync::Arc;

use derive_more::{IsVariant, Unwrap};

use super::common as c;
use super::common::Catchable;
use super::public as p;

use c::Symbol;

use crate::bytecode::OpCodes;
use crate::error::*;
use crate::vm::{Env, VM};

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

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

pub trait ToPublic {
    fn to_public(self, vm: &VM) -> p::Value;
}

#[derive(Debug, Clone)]
pub struct Thunk {
    pub thunk: RefCell<_Thunk>,
    pub env: RefCell<Option<Arc<Env>>>,
}

#[derive(Debug, IsVariant, Unwrap, Clone)]
pub enum _Thunk {
    Code(OpCodes),
    SuspendedFrom(*const Thunk),
    Value(Box<Value>),
}

impl Thunk {
    pub fn new(opcodes: OpCodes) -> Thunk {
        Thunk {
            thunk: RefCell::new(_Thunk::Code(opcodes)),
            env: RefCell::new(None),
        }
    }

    pub fn unwrap_code(&self) -> OpCodes {
        self.thunk.borrow().clone().unwrap_code()
    }

    pub fn capture(&self, env: Arc<Env>) {
        *self.env.borrow_mut() = Some(env);
    }

    pub fn force(&self, vm: &VM) -> Result<Value> {
        match &*self.thunk.borrow() {
            _Thunk::Value(value) => return Ok(value.as_ref().clone()),
            _Thunk::SuspendedFrom(from) => {
                return Err(Error::EvalError(format!(
                    "thunk {:p} already suspended from {from:p} (infinite recursion encountered)",
                    self as *const Thunk
                )));
            }
            _Thunk::Code(_) => (),
        }
        let opcodes = std::mem::replace(
            &mut *self.thunk.borrow_mut(),
            _Thunk::SuspendedFrom(self as *const Thunk),
        )
        .unwrap_code();
        let value = vm.eval(opcodes, self.env.borrow().clone().unwrap())?;
        let _ = std::mem::replace(
            &mut *self.thunk.borrow_mut(),
            _Thunk::Value(value.clone().into()),
        );
        Ok(value)
    }

    pub fn value(&self) -> Option<Value> {
        match &*self.thunk.borrow() {
            _Thunk::Value(value) => Some(value.as_ref().clone()),
            _ => None,
        }
    }
}

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

#[derive(Debug, IsVariant, Unwrap, Clone, PartialEq)]
pub enum Value {
    Const(Const),
    Thunk(Thunk),
    AttrSet(AttrSet),
    RecAttrSet(RecAttrSet),
    List(List),
    Catchable(c::Catchable),
    PrimOp(PrimOp),
    PartialPrimOp(PartialPrimOp),
    Func(Func),
}

#[derive(Debug, IsVariant, Unwrap, Clone, PartialEq)]
pub enum ValueAsRef<'a> {
    Const(&'a Const),
    Thunk(&'a Thunk),
    AttrSet(&'a AttrSet),
    RecAttrSet(&'a RecAttrSet),
    List(&'a List),
    Catchable(&'a c::Catchable),
    PrimOp(&'a PrimOp),
    PartialPrimOp(&'a PartialPrimOp),
    Func(&'a Func),
}

#[derive(Debug, IsVariant, Unwrap, PartialEq)]
pub enum ValueAsMut<'a> {
    Const(&'a mut Const),
    Thunk(&'a mut Thunk),
    AttrSet(&'a mut AttrSet),
    RecAttrSet(&'a mut RecAttrSet),
    List(&'a mut List),
    Catchable(&'a mut c::Catchable),
    PrimOp(&'a mut PrimOp),
    PartialPrimOp(&'a mut PartialPrimOp),
    Func(&'a mut Func),
}

impl Value {
    pub fn as_ref(&self) -> ValueAsRef<'_> {
        use Value::*;
        use ValueAsRef as R;
        match self {
            Const(x) => R::Const(x),
            Thunk(x) => R::Thunk(x),
            AttrSet(x) => R::AttrSet(x),
            RecAttrSet(x) => R::RecAttrSet(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(&mut self) -> ValueAsMut<'_> {
        use Value::*;
        use ValueAsMut as M;
        match self {
            Const(x) => M::Const(x),
            Thunk(x) => M::Thunk(x),
            AttrSet(x) => M::AttrSet(x),
            RecAttrSet(x) => M::RecAttrSet(x),
            List(x) => M::List(x),
            Catchable(x) => M::Catchable(x),
            PrimOp(x) => M::PrimOp(x),
            PartialPrimOp(x) => M::PartialPrimOp(x),
            Func(x) => M::Func(x),
        }
    }
}

use Value::Const as VmConst;
impl Value {
    pub fn typename(&self) -> &'static str {
        use Value::*;
        match self {
            Const(_) => todo!(),
            Thunk(_) => "thunk",
            AttrSet(_) => "set",
            RecAttrSet(_) => "set",
            List(_) => "list",
            Catchable(_) => todo!(),
            PrimOp(_) => "lambda",
            PartialPrimOp(_) => "lambda",
            Func(_) => "lambda",
        }
    }

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

    pub fn call(self, vm: &VM, args: Vec<Value>) -> Result<Value> {
        use Value::*;
        Ok(match self {
            PrimOp(func) => func.call(vm, args),
            PartialPrimOp(func) => func.call(vm, args),
            mut func @ Value::Func(_) => {
                let mut iter = args.into_iter();
                while let Some(arg) = iter.next() {
                    func = match func {
                        PrimOp(func) => {
                            return Ok(func.call(vm, [arg].into_iter().chain(iter).collect()));
                        }
                        PartialPrimOp(func) => {
                            return Ok(func.call(vm, [arg].into_iter().chain(iter).collect()));
                        }
                        Func(func) => func.call(vm, arg)?,
                        _ => todo!(),
                    }
                }
                func
            }
            x @ Catchable(_) => x,
            _ => todo!(),
        })
    }

    pub fn not(self) -> Value {
        use Const::*;
        match self {
            VmConst(Bool(bool)) => VmConst(Bool(!bool)),
            x @ Value::Catchable(_) => x,
            _ => todo!(),
        }
    }

    pub fn and(self, other: Value) -> Value {
        use Const::*;
        match (self, other) {
            (VmConst(Bool(a)), VmConst(Bool(b))) => VmConst(Bool(a && b)),
            _ => todo!(),
        }
    }

    pub fn or(self, other: Value) -> Value {
        use Const::*;
        match (self, other) {
            (VmConst(Bool(a)), VmConst(Bool(b))) => VmConst(Bool(a || b)),
            _ => todo!(),
        }
    }

    pub fn eq(self, other: Value) -> Value {
        use Const::Bool;
        VmConst(Bool(self == other))
    }

    pub fn lt(self, other: Value) -> Value {
        use Const::*;
        VmConst(Bool(match (self, other) {
            (VmConst(Int(a)), VmConst(Int(b))) => a < b,
            (VmConst(Int(a)), VmConst(Float(b))) => (a as f64) < b,
            (VmConst(Float(a)), VmConst(Int(b))) => a < b as f64,
            (VmConst(Float(a)), VmConst(Float(b))) => a < b,
            (VmConst(String(a)), VmConst(String(b))) => a < b,
            _ => todo!(),
        }))
    }

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

    pub fn add(self, other: Value) -> Value {
        use Const::*;
        match (self, other) {
            (VmConst(Int(a)), VmConst(Int(b))) => VmConst(Int(a + b)),
            (VmConst(Int(a)), VmConst(Float(b))) => VmConst(Float(a as f64 + b)),
            (VmConst(Float(a)), VmConst(Int(b))) => VmConst(Float(a + b as f64)),
            (VmConst(Float(a)), VmConst(Float(b))) => VmConst(Float(a + b)),
            (VmConst(String(a)), VmConst(String(b))) => {
                let mut string = a.clone();
                string.push_str(b.as_str());
                VmConst(String(string))
            }
            _ => todo!(),
        }
    }

    pub fn mul(self, other: Value) -> Value {
        use Const::*;
        match (self, other) {
            (VmConst(Int(a)), VmConst(Int(b))) => VmConst(Int(a * b)),
            (VmConst(Int(a)), VmConst(Float(b))) => VmConst(Float(a as f64 * b)),
            (VmConst(Float(a)), VmConst(Int(b))) => VmConst(Float(a * b as f64)),
            (VmConst(Float(a)), VmConst(Float(b))) => VmConst(Float(a * b)),
            _ => todo!(),
        }
    }

    pub fn div(self, other: Value) -> Value {
        use Const::*;
        match (self, other) {
            (_, VmConst(Int(0))) => todo!(),
            (_, VmConst(Float(0.))) => todo!(),
            (VmConst(Int(a)), VmConst(Int(b))) => VmConst(Int(a / b)),
            (VmConst(Int(a)), VmConst(Float(b))) => VmConst(Float(a as f64 / b)),
            (VmConst(Float(a)), VmConst(Int(b))) => VmConst(Float(a / b as f64)),
            (VmConst(Float(a)), VmConst(Float(b))) => VmConst(Float(a / b)),
            _ => todo!(),
        }
    }

    pub fn concat_string(&mut self, mut other: Value) -> &mut Self {
        if let (VmConst(Const::String(a)), VmConst(Const::String(b))) =
            (self.coerce_to_string(), other.coerce_to_string())
        {
            a.push_str(b.as_str());
        } else {
            todo!()
        }
        self
    }

    pub fn push(&mut self, elem: Value) -> &mut Self {
        if let Value::List(list) = self {
            list.push(elem);
        } else {
            todo!()
        }
        self
    }

    pub fn concat(self, other: Value) -> Value {
        if let (Value::List(a), Value::List(b)) = (self, other) {
            Value::List(a.concat(b))
        } else {
            todo!()
        }
    }

    pub fn push_attr(&mut self, sym: Symbol, val: Value) -> &mut Self {
        if let Value::AttrSet(attrs) = self {
            attrs.push_attr(sym, val)
        } else if let Value::RecAttrSet(attrs) = self {
            attrs.push_attr(sym, val)
        } else {
            todo!()
        }
        self
    }

    pub fn update(self, other: Value) -> Value {
        match (self, other) {
            (Value::AttrSet(a), Value::AttrSet(b)) => Value::AttrSet(a.update(b)),
            (Value::RecAttrSet(a), Value::AttrSet(b)) => Value::AttrSet(a.update_normal(b)),
            (Value::AttrSet(a), Value::RecAttrSet(b)) => Value::AttrSet(a.update_rec(b)),
            _ => todo!(),
        }
    }

    pub fn select(&mut self, sym: Symbol) -> &mut Self {
        if let Value::AttrSet(attrs) = self {
            let val = attrs.select(sym.clone()).unwrap_or_else(|| {
                Value::Catchable(c::Catchable::new(Some(format!("{sym:?} not found"))))
            });
            *self = val;
        } else if let Value::RecAttrSet(attrs) = self {
            let val = attrs.select(sym.clone()).unwrap_or_else(|| {
                Value::Catchable(c::Catchable::new(Some(format!("{sym:?} not found"))))
            });
            *self = val;
        } else {
            *self = Value::Catchable(Catchable::new(Some(format!(
                "cannot select from {:?}",
                self.typename()
            ))))
        }
        self
    }

    pub fn select_with_default(&mut self, sym: Symbol, default: Value) -> &mut Self {
        if let Value::AttrSet(attrs) = self {
            let val = attrs.select(sym).unwrap_or(default);
            *self = val;
        } else if let Value::RecAttrSet(attrs) = self {
            let val = attrs.select(sym).unwrap_or(default);
            *self = val;
        } else {
            todo!()
        }
        self
    }

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

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

    pub fn force(&mut self, vm: &VM) -> Result<&mut Self> {
        if let Value::Thunk(thunk) = self {
            let value = thunk.force(vm)?;
            *self = value
        }
        Ok(self)
    }

    pub fn force_deep(&mut self, vm: &VM) -> Result<&mut Self> {
        match self {
            Value::Thunk(thunk) => {
                let mut value = thunk.force(vm)?;
                value.force_deep(vm)?;
                *self = value;
            }
            Value::List(list) => list.force_deep(vm)?,
            Value::AttrSet(attrs) => attrs.force_deep(vm)?,
            Value::RecAttrSet(attrs) => attrs.force_deep(vm)?,
            _ => (),
        }
        Ok(self)
    }
}

impl ToPublic for Value {
    fn to_public(self, vm: &VM) -> p::Value {
        match self {
            Value::AttrSet(attrs) => attrs.to_public(vm),
            Value::RecAttrSet(attrs) => attrs.to_public(vm),
            Value::List(list) => list.to_public(vm),
            Value::Catchable(catchable) => p::Value::Catchable(catchable),
            Value::Const(cnst) => p::Value::Const(cnst.into()),
            Value::Thunk(_) => p::Value::Thunk,
            Value::PrimOp(primop) => p::Value::PrimOp(primop.name),
            Value::PartialPrimOp(primop) => p::Value::PartialPrimOp(primop.name),
            Value::Func(_) => p::Value::Func,
        }
    }
}