use std::{fmt::Debug, marker::PhantomData};
use crate::index::Idx;
#[derive(Clone, Copy)]
pub struct Embedding<I: Idx, J: Idx> {
em_start: u32,
#[cfg(debug_assertions)]
em_size: u32,
_marker: PhantomData<(I, J)>,
}
impl<I: Idx, J: Idx> Embedding<I, J> {
pub fn get(self, index: J) -> I {
#[cfg(debug_assertions)]
{
assert!(index.index() < self.em_size as u64);
}
Idx::new(self.em_start as u64 + index.index())
}
}
impl<I: Idx> Debug for Alloc<I> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let next = <I as Idx>::new(self.cur_head as u64);
f.debug_struct("Alloc").field("next", &next).finish()
}
}
pub struct Alloc<I: Idx> {
cur_head: u32,
_marker: PhantomData<I>,
}
impl<I: Idx> Alloc<I> {
pub fn new() -> Alloc<I> {
Alloc {
cur_head: 0,
_marker: PhantomData,
}
}
pub fn new_at(min: usize) -> Alloc<I> {
Alloc {
cur_head: min as u32,
_marker: PhantomData,
}
}
pub fn alloc(&mut self) -> I {
self.cur_head += 1;
Idx::new((self.cur_head - 1) as u64)
}
pub fn alloc_embedding<J: Idx>(&mut self, len: usize) -> Embedding<I, J> {
let start = self.cur_head;
self.cur_head += len as u32;
Embedding {
em_start: start,
#[cfg(debug_assertions)]
em_size: len as u32,
_marker: PhantomData,
}
}
}