test.rs

use std::fmt::Debug;
use std::fmt;

#[derive(Debug, Copy, Clone, PartialEq)]
pub struct WFuncID(u32);
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct WTableID(u32);

#[allow(non_snake_case)]
#[derive(Debug, Copy, Clone, PartialEq)]
pub enum WASMType<HostTy: Debug + Clone + Send + Sync + PartialEq> {
  I32,
  I64,
  F32,
  F64,
  V128,
  NullRef,
  HostRef,
  FuncRef,
  TableRef,
  // TODO: figure out how to get rid of this
  #[doc(hidden)]
  __PhantomHostRef(std::marker::PhantomData<HostTy>),
}

#[allow(non_snake_case)]
#[derive(Debug, Clone, PartialEq)]
pub enum WASMValue<HostTy: Debug + Clone + Send + Sync + PartialEq> {
  I32(i32),
  I64(i64),
  F32(f32),
  F64(f64),
  V128(u128),
  NullRef(),
  HostRef(HostTy),
  FuncRef(WFuncID), // u32 should be plenty.
  TableRef(WTableID),
}

#[allow(non_snake_case)]
#[derive(Copy, Clone)]
#[repr(C)]
union WSTNumV {
  I32: i32,
  I64: i64,
  F32: f32,
  F64: f64,
}

impl WSTNumV {
  unsafe fn cast<HostTy: Debug + Clone + Send + Sync + PartialEq>(self, ty: WASMType<HostTy>) -> Option<WASMValue<HostTy>> {
    match ty {
      WASMType::I32 => {
        Some(WASMValue::I32::<HostTy>(self.I32))
      },
      WASMType::I64 => {
        Some(WASMValue::I64::<HostTy>(self.I64))
      },
      WASMType::F32 => {
        Some(WASMValue::F32::<HostTy>(self.F32))
      },
      WASMType::F64 => {
        Some(WASMValue::F64::<HostTy>(self.F64))
      },
      _ => {
        None
      },
    }
  }
}

impl fmt::Debug for WSTNumV {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    // Is "sound". We're using this to generate debug data about a potentially invalid value, so the value's data being misinterpreted is fine.
    unsafe {
      write!(f, "WSTNumV {{ i32: {}, i64: {}, f32: {:e}, f64: {:e}, hex: {:x} }}", self.I32, self.I64, self.F32, self.F64, self.I64)
    }
  }
}

#[derive(Debug)]
pub struct WASMStack<HostTy: Debug + Clone + Send + Sync + PartialEq> {
  numstack: Vec<WSTNumV>,
  typestack: Vec<WASMType<HostTy>>,
  vecstack: Vec<u128>,
  hoststack: Vec<HostTy>, // do NOT conjoin hoststack with anything. Host type is unpredictable.
  funcstack: Vec<WFuncID>, // Maybe conjoin funcstack and tablestack?
  tablestack: Vec<WTableID>,
  max: usize,
}

#[derive(Debug)]
pub enum WStackError {
  StackOverflow,
  StackUnderflow,
  WrongType,
}

impl std::fmt::Display for WStackError {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    match self {
      WStackError::StackOverflow => write!(f, "Stack Overflow"),
      WStackError::StackUnderflow => write!(f, "Stack Underflow"),
      WStackError::WrongType => write!(f, "Wrong type"),
    }
  }
}

impl std::error::Error for WStackError {
  fn description(&self) -> &str {
    "Stack overflow/underflow"
  }
}

pub type SResult<T> = std::result::Result<T, WStackError>;

impl<HostTy: Debug + Clone + Send + Sync + PartialEq> WASMStack<HostTy> {
  pub fn new(max: usize) -> WASMStack<HostTy> {
    WASMStack {
      numstack: vec![],
      typestack: vec![],
      vecstack: vec![],
      hoststack: vec![],
      funcstack: vec![],
      tablestack: vec![],
      max,
    }
  }

  fn check_overflow(&mut self) -> SResult<()> {
    if self.typestack.len() >= self.max {
      return Err(WStackError::StackOverflow);
    }
    Ok(())
  }

  pub fn push(&mut self, v: WASMValue<HostTy>) -> SResult<()> {
    match v {
      WASMValue::I32(n) => self.push_i32(n)?,
      WASMValue::I64(n) => self.push_i64(n)?,
      WASMValue::F32(n) => self.push_f32(n)?,
      WASMValue::F64(n) => self.push_f64(n)?,
      WASMValue::V128(n) => self.push_v128(n)?,
      WASMValue::NullRef() => self.push_nullref()?,
      WASMValue::HostRef(n) => self.push_hostref(n)?,
      WASMValue::FuncRef(n) => self.push_funcref(n)?,
      WASMValue::TableRef(n) => self.push_tableref(n)?,
    }
    Ok(())
  }

  pub fn pop(&mut self) -> SResult<WASMValue<HostTy>> {
    let t = self.typestack.pop();

    match t {
      Some(WASMType::I32) |
      Some(WASMType::I64) |
      Some(WASMType::F32) |
      Some(WASMType::F64) => {
        Ok (
        unsafe {
          self.numstack.pop()
          .expect("Internal Stack Underflow! This is a memory corruption bug! ):")
          .cast(t.unwrap()).unwrap()
        }
        )
      },
      Some(WASMType::V128) => {
        let v = self.vecstack.pop()
        .expect("Internal Stack Underflow! This is a memory corruption bug! ):");
        Ok(WASMValue::V128(v))
      },
      Some(WASMType::NullRef) => {
        Ok(WASMValue::NullRef())
      },
      Some(WASMType::HostRef) => {
        let v = self.hoststack.pop()
        .expect("Internal Stack Underflow! This is a memory corruption bug! ):");
        Ok(WASMValue::HostRef(v))
      },
      Some(WASMType::FuncRef) => {
        let v = self.funcstack.pop()
        .expect("Internal Stack Underflow! This is a memory corruption bug! ):");
        Ok(WASMValue::FuncRef(v))
      },
      Some(WASMType::TableRef) => {
        let v = self.tablestack.pop()
        .expect("Internal Stack Underflow! This is a memory corruption bug! ):");
        Ok(WASMValue::TableRef(v))
      },
      Some(WASMType::__PhantomHostRef(_)) => {
        unreachable!("PhantomType cannot be constructed.")
      }
      None => {
        Err(WStackError::StackUnderflow)
      }
    }
  }
  
  pub fn push_i32(&mut self, n: i32) -> SResult<()> {
    self.check_overflow()?;
    self.typestack.push(WASMType::I32::<HostTy>);
    self.numstack.push(WSTNumV {I32: n});
    Ok(())
  }

  pub fn push_i64(&mut self, n: i64) -> SResult<()> {
    self.check_overflow()?;
    self.typestack.push(WASMType::I64::<HostTy>);
    self.numstack.push(WSTNumV {I64: n});
    Ok(())
  }

  pub fn push_f32(&mut self, n: f32) -> SResult<()> {
    self.check_overflow()?;
    self.typestack.push(WASMType::F32::<HostTy>);
    self.numstack.push(WSTNumV {F32: n});
    Ok(())
  }

  pub fn push_f64(&mut self, n: f64) -> SResult<()> {
    self.check_overflow()?;
    self.typestack.push(WASMType::F64::<HostTy>);
    self.numstack.push(WSTNumV {F64: n});
    Ok(())
  }

  pub fn push_v128(&mut self, n: u128) -> SResult<()> {
    self.check_overflow()?;
    self.typestack.push(WASMType::V128::<HostTy>);
    self.vecstack.push(n);
    Ok(())
  }

  pub fn push_nullref(&mut self) -> SResult<()> {
    self.check_overflow()?;
    self.typestack.push(WASMType::NullRef::<HostTy>);
    Ok(())
  }

  pub fn push_hostref(&mut self, n: HostTy) -> SResult<()> {
    self.check_overflow()?;
    self.typestack.push(WASMType::HostRef::<HostTy>);
    self.hoststack.push(n);
    Ok(())
  }

  pub fn push_funcref(&mut self, n: WFuncID) -> SResult<()> {
    self.check_overflow()?;
    self.typestack.push(WASMType::FuncRef::<HostTy>);
    self.funcstack.push(n);
    Ok(())
  }

  pub fn push_tableref(&mut self, n: WTableID) -> SResult<()> {
    self.check_overflow()?;
    self.typestack.push(WASMType::TableRef::<HostTy>);
    self.tablestack.push(n);
    Ok(())
  }

  pub fn pop_i32(&mut self) -> SResult<i32> {
    let t = self.typestack.pop();
    match t {
      Some(WASMType::I32) => {
        let v = self.numstack.pop()
          .expect("Internal Stack Underflow! This is a memory corruption bug! ):");
        unsafe {Ok(v.I32)}
      }
      Some(_) => {Err(WStackError::WrongType) }
      None => {Err(WStackError::StackUnderflow)}
    }
  }

  pub fn pop_i64(&mut self) -> SResult<i64> {
    let t = self.typestack.pop();
    match t {
      Some(WASMType::I64) => {
        let v = self.numstack.pop()
          .expect("Internal Stack Underflow! This is a memory corruption bug! ):");
        unsafe {Ok(v.I64)}
      }
      Some(_) => {Err(WStackError::WrongType) }
      None => {Err(WStackError::StackUnderflow)}
    }
  }

  pub fn pop_f32(&mut self) -> SResult<f32> {
    let t = self.typestack.pop();
    match t {
      Some(WASMType::F32) => {
        let v = self.numstack.pop()
          .expect("Internal Stack Underflow! This is a memory corruption bug! ):");
        unsafe {Ok(v.F32)}
      }
      Some(_) => {Err(WStackError::WrongType) }
      None => {Err(WStackError::StackUnderflow)}
    }
  }

  pub fn pop_f64(&mut self) -> SResult<f64> {
    let t = self.typestack.pop();
    match t {
      Some(WASMType::F64) => {
        let v = self.numstack.pop()
          .expect("Internal Stack Underflow! This is a memory corruption bug! ):");
        unsafe {Ok(v.F64)}
      }
      Some(_) => {Err(WStackError::WrongType) }
      None => {Err(WStackError::StackUnderflow)}
    }
  }

  pub fn pop_v128(&mut self) -> SResult<u128> {
    let t = self.typestack.pop();
    match t {
      Some(WASMType::V128) => {
        let v = self.vecstack.pop()
          .expect("Internal Stack Underflow! This is a memory corruption bug! ):");
        Ok(v)
      }
      Some(_) => {Err(WStackError::WrongType) }
      None => {Err(WStackError::StackUnderflow)}
    }
  }

  pub fn pop_nullref(&mut self) -> SResult<()> {
    let t = self.typestack.pop();
    match t {
      Some(WASMType::NullRef) => {
        Ok(())
      }
      Some(_) => {Err(WStackError::WrongType) }
      None => {Err(WStackError::StackUnderflow)}
    }
  }

  pub fn pop_hostref(&mut self) -> SResult<HostTy> {
    let t = self.typestack.pop();
    match t {
      Some(WASMType::HostRef) => {
        let v = self.hoststack.pop()
          .expect("Internal Stack Underflow! This is a memory corruption bug! ):");
        Ok(v)
      }
      Some(_) => {Err(WStackError::WrongType) }
      None => {Err(WStackError::StackUnderflow)}
    }
  }

  pub fn pop_funcref(&mut self) -> SResult<WFuncID> {
    let t = self.typestack.pop();
    match t {
      Some(WASMType::FuncRef) => {
        let v = self.funcstack.pop()
          .expect("Internal Stack Underflow! This is a memory corruption bug! ):");
        Ok(v)
      }
      Some(_) => {Err(WStackError::WrongType) }
      None => {Err(WStackError::StackUnderflow)}
    }
  }

  pub fn pop_tableref(&mut self) -> SResult<WTableID> {
    let t = self.typestack.pop();
    match t {
      Some(WASMType::TableRef) => {
        let v = self.tablestack.pop()
          .expect("Internal Stack Underflow! This is a memory corruption bug! ):");
        Ok(v)
      }
      Some(_) => {Err(WStackError::WrongType) }
      None => {Err(WStackError::StackUnderflow)}
    }
  }

  // Copy back the type of a value `depth` units deep.
  pub fn pick_type(&self, depth: usize) -> SResult<WASMType<HostTy>> {
    let top = self.typestack.len()-1;

    if let Some(t) = self.typestack.get(top - depth) {
      return Ok(t.clone());
    }
    return Err(WStackError::StackUnderflow);
  }

  pub fn dump_stack(mut self) -> Vec<WASMValue<HostTy>> {
    let mut out = vec![];
    while let Ok(v) = self.pop() {
      out.push(v);
    }

    out
  }

  pub fn map_stack<'a>(&'a self) -> &'a Vec<WASMType<HostTy>> {
    &self.typestack
  }
}