goosedump 0.11.2

// SPDX-License-Identifier: LGPL-2.1-or-later
// Copyright (C) Jarkko Sakkinen 2026

//! Minimal validated GGUF v3 reader for the fixed inference engines.

use std::borrow::Cow;
use std::collections::HashMap;
use std::fs::File;
use std::path::Path;
use std::str;

use anyhow::{Context as _, Result, anyhow, bail};
use memmap2::{Mmap, MmapOptions};

const MAGIC: &[u8; 4] = b"GGUF";
const VERSION: u32 = 3;
const DEFAULT_ALIGNMENT: usize = 32;
const MAX_DIMENSIONS: usize = 4;
const Q8_BLOCK_VALUES: usize = 32;
const Q8_BLOCK_BYTES: usize = 34;
const MXFP4_BLOCK_VALUES: usize = 32;
const MXFP4_BLOCK_BYTES: usize = 17;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(super) enum TensorType {
    F32,
    Bf16,
    Q8_0,
    Mxfp4,
}

#[derive(Clone, Copy, Debug)]
pub(super) struct Tensor<'a> {
    dimensions: &'a [usize],
    kind: TensorType,
    data: &'a [u8],
}

impl<'a> Tensor<'a> {
    pub(super) fn dimensions(&self) -> &[usize] {
        self.dimensions
    }

    pub(super) fn tensor_type(&self) -> TensorType {
        self.kind
    }

    pub(super) fn f32_slice(&self) -> Result<&'a [f32]> {
        if self.kind != TensorType::F32 {
            bail!("tensor is {:?}, not F32", self.kind);
        }
        if cfg!(target_endian = "big") {
            bail!("F32 tensor views require a little-endian host");
        }
        if !(self.data.as_ptr() as usize).is_multiple_of(align_of::<f32>()) {
            bail!("F32 tensor data is not aligned");
        }
        // Every bit pattern is a valid f32 and the parser checked the byte count.
        // The alignment guard above makes `cast_slice` (which otherwise panics on
        // misalignment) infallible here, with no `unsafe` at the call site.
        Ok(bytemuck::cast_slice(self.data))
    }

    pub(super) fn row_count(&self) -> usize {
        self.dimensions.get(1).copied().unwrap_or(1)
    }

    pub(super) fn f32_row(&self, row: usize) -> Result<&'a [f32]> {
        let width = self.dimensions[0];
        if row >= self.row_count() {
            bail!("F32 row {row} is out of range");
        }
        let values = self.f32_slice()?;
        Ok(&values[row * width..(row + 1) * width])
    }

    pub(super) fn bf16_row(&self, row: usize) -> Result<&'a [u8]> {
        if self.kind != TensorType::Bf16 {
            bail!("tensor is {:?}, not BF16", self.kind);
        }
        if row >= self.row_count() {
            bail!("BF16 row {row} is out of range");
        }
        let row_size = self.dimensions[0] * size_of::<u16>();
        Ok(&self.data[row * row_size..(row + 1) * row_size])
    }

    pub(super) fn q8_row(&self, row: usize) -> Result<&'a [u8]> {
        if self.kind != TensorType::Q8_0 {
            bail!("tensor is {:?}, not Q8_0", self.kind);
        }
        if row >= self.row_count() {
            bail!("Q8_0 row {row} is out of range");
        }
        let row_size = self.dimensions[0] / Q8_BLOCK_VALUES * Q8_BLOCK_BYTES;
        Ok(&self.data[row * row_size..(row + 1) * row_size])
    }

    pub(super) fn mxfp4_row(&self, row: usize) -> Result<&'a [u8]> {
        if self.kind != TensorType::Mxfp4 {
            bail!("tensor is {:?}, not MXFP4", self.kind);
        }
        if row >= self.row_count() {
            bail!("MXFP4 row {row} is out of range");
        }
        let row_size = self.dimensions[0] / MXFP4_BLOCK_VALUES * MXFP4_BLOCK_BYTES;
        Ok(&self.data[row * row_size..(row + 1) * row_size])
    }

    pub(super) fn matrix_slice(&self, index: usize) -> Result<Self> {
        let &[_, _, count] = self.dimensions else {
            bail!(
                "tensor has dimensions {:?}, expected three",
                self.dimensions
            );
        };
        if index >= count {
            bail!("tensor matrix {index} is out of range");
        }
        let byte_len = self.data.len() / count;
        Ok(Self {
            dimensions: &self.dimensions[..2],
            kind: self.kind,
            data: &self.data[index * byte_len..(index + 1) * byte_len],
        })
    }
}

pub(super) struct Gguf {
    bytes: Mmap,
    data_start: usize,
    metadata: HashMap<String, MetadataValue>,
    tensors: HashMap<String, TensorInfo>,
}

impl Gguf {
    pub(super) fn load(path: impl AsRef<Path>) -> Result<Self> {
        let path = path.as_ref();
        let file = File::open(path).with_context(|| format!("open {}", path.display()))?;
        // The mapping is read-only and remains owned by Gguf for every tensor view.
        let bytes = unsafe { MmapOptions::new().map(&file) }
            .with_context(|| format!("map {}", path.display()))?;
        Self::parse(bytes).with_context(|| format!("parse {}", path.display()))
    }

    fn parse(bytes: Mmap) -> Result<Self> {
        let mut reader = Reader::new(&bytes);
        if reader.take(MAGIC.len())? != MAGIC {
            bail!("invalid GGUF magic");
        }
        let version = reader.u32()?;
        if version != VERSION {
            bail!("unsupported GGUF version {version}, expected {VERSION}");
        }
        let tensor_count = reader.count("tensor count")?;
        let metadata_count = reader.count("metadata count")?;
        if metadata_count > reader.remaining() / 13 {
            bail!("metadata count exceeds the file size");
        }

        let mut metadata = HashMap::new();
        for _ in 0..metadata_count {
            let name = reader.string("metadata key")?;
            validate_name(&name, "metadata key")?;
            let value_type = reader.u32()?;
            let value = MetadataValue::read(&mut reader, value_type)
                .with_context(|| format!("read metadata `{name}`"))?;
            if metadata.insert(name.clone(), value).is_some() {
                bail!("duplicate metadata key `{name}`");
            }
        }
        validate_architecture(&metadata)?;
        let alignment = metadata_alignment(&metadata)?;
        if tensor_count > reader.remaining() / 32 {
            bail!("tensor count exceeds the file size");
        }

        let mut tensors = HashMap::new();
        for _ in 0..tensor_count {
            let name = reader.string("tensor name")?;
            validate_name(&name, "tensor name")?;
            let dimension_count =
                usize::try_from(reader.u32()?).context("tensor dimension count")?;
            if !(1..=MAX_DIMENSIONS).contains(&dimension_count) {
                bail!("tensor `{name}` has invalid dimension count {dimension_count}");
            }
            let mut dimensions = Vec::with_capacity(dimension_count);
            for _ in 0..dimension_count {
                let dimension = reader.usize("tensor dimension")?;
                if dimension == 0 {
                    bail!("tensor `{name}` has a zero dimension");
                }
                dimensions.push(dimension);
            }
            let kind = match reader.u32()? {
                0 => TensorType::F32,
                8 => TensorType::Q8_0,
                30 => TensorType::Bf16,
                39 => TensorType::Mxfp4,
                value => bail!("tensor `{name}` uses unsupported GGML type {value}"),
            };
            let offset = reader.usize("tensor offset")?;
            if offset % alignment != 0 {
                bail!("tensor `{name}` offset is not {alignment}-byte aligned");
            }
            let byte_len = tensor_byte_len(&name, &dimensions, kind)?;
            let info = TensorInfo {
                dimensions,
                kind,
                offset,
                byte_len,
            };
            if tensors.insert(name.clone(), info).is_some() {
                bail!("duplicate tensor name `{name}`");
            }
        }
        let data_start = align_up(reader.position(), alignment)?;
        validate_ranges(&bytes, data_start, &tensors)?;
        Ok(Self {
            bytes,
            data_start,
            metadata,
            tensors,
        })
    }

    pub(super) fn architecture(&self) -> &str {
        match self.metadata.get("general.architecture") {
            Some(MetadataValue::String(value)) => value,
            _ => unreachable!(),
        }
    }

    pub(super) fn u32(&self, key: &str) -> Result<u32> {
        match self.value(key)? {
            MetadataValue::U32(value) => Ok(*value),
            value => type_error(key, "u32", value),
        }
    }

    pub(super) fn f32(&self, key: &str) -> Result<f32> {
        match self.value(key)? {
            MetadataValue::F32(value) => Ok(*value),
            value => type_error(key, "f32", value),
        }
    }

    pub(super) fn string(&self, key: &str) -> Result<&str> {
        match self.value(key)? {
            MetadataValue::String(value) => Ok(value),
            value => type_error(key, "string", value),
        }
    }

    pub(super) fn strings(&self, key: &str) -> Result<&[String]> {
        match self.value(key)? {
            MetadataValue::Array(MetadataArray::String(values)) => Ok(values),
            value => type_error(key, "string array", value),
        }
    }

    pub(super) fn i32s(&self, key: &str) -> Result<Cow<'_, [i32]>> {
        match self.value(key)? {
            MetadataValue::Array(MetadataArray::I32(values)) => Ok(Cow::Borrowed(values)),
            MetadataValue::Array(MetadataArray::U32(values)) => values
                .iter()
                .copied()
                .map(|value| i32::try_from(value).context("token type exceeds i32"))
                .collect::<Result<Vec<_>>>()
                .map(Cow::Owned),
            value => type_error(key, "i32 array", value),
        }
    }

    pub(super) fn tensor(&self, name: &str) -> Result<Tensor<'_>> {
        let info = self
            .tensors
            .get(name)
            .ok_or_else(|| anyhow!("missing tensor `{name}`"))?;
        let start = self.data_start + info.offset;
        Ok(Tensor {
            dimensions: &info.dimensions,
            kind: info.kind,
            data: &self.bytes[start..start + info.byte_len],
        })
    }

    fn value(&self, key: &str) -> Result<&MetadataValue> {
        self.metadata
            .get(key)
            .ok_or_else(|| anyhow!("missing GGUF metadata `{key}`"))
    }
}

struct TensorInfo {
    dimensions: Vec<usize>,
    kind: TensorType,
    offset: usize,
    byte_len: usize,
}

#[allow(dead_code)]
enum MetadataValue {
    U8(u8),
    I8(i8),
    U16(u16),
    I16(i16),
    U32(u32),
    I32(i32),
    F32(f32),
    Bool(bool),
    String(String),
    Array(MetadataArray),
    U64(u64),
    I64(i64),
    F64(f64),
}

impl MetadataValue {
    fn read(reader: &mut Reader<'_>, kind: u32) -> Result<Self> {
        match kind {
            0 => Ok(Self::U8(reader.u8()?)),
            1 => Ok(Self::I8(reader.i8()?)),
            2 => Ok(Self::U16(reader.u16()?)),
            3 => Ok(Self::I16(reader.i16()?)),
            4 => Ok(Self::U32(reader.u32()?)),
            5 => Ok(Self::I32(reader.i32()?)),
            6 => Ok(Self::F32(reader.f32()?)),
            7 => Ok(Self::Bool(reader.bool()?)),
            8 => Ok(Self::String(reader.string("metadata string")?)),
            9 => Ok(Self::Array(MetadataArray::read(reader)?)),
            10 => Ok(Self::U64(reader.u64()?)),
            11 => Ok(Self::I64(reader.i64()?)),
            12 => Ok(Self::F64(reader.f64()?)),
            _ => bail!("unknown GGUF metadata type {kind}"),
        }
    }

    fn type_name(&self) -> &'static str {
        match self {
            Self::U8(_) => "u8",
            Self::I8(_) => "i8",
            Self::U16(_) => "u16",
            Self::I16(_) => "i16",
            Self::U32(_) => "u32",
            Self::I32(_) => "i32",
            Self::F32(_) => "f32",
            Self::Bool(_) => "bool",
            Self::String(_) => "string",
            Self::Array(value) => value.type_name(),
            Self::U64(_) => "u64",
            Self::I64(_) => "i64",
            Self::F64(_) => "f64",
        }
    }
}

#[allow(dead_code)]
enum MetadataArray {
    U8(Vec<u8>),
    I8(Vec<i8>),
    U16(Vec<u16>),
    I16(Vec<i16>),
    U32(Vec<u32>),
    I32(Vec<i32>),
    F32(Vec<f32>),
    Bool(Vec<bool>),
    String(Vec<String>),
    U64(Vec<u64>),
    I64(Vec<i64>),
    F64(Vec<f64>),
}

impl MetadataArray {
    fn read(reader: &mut Reader<'_>) -> Result<Self> {
        let kind = reader.u32()?;
        if kind == 9 || kind > 12 {
            bail!("invalid GGUF metadata array type {kind}");
        }
        let count = reader.count("metadata array length")?;
        let minimum = match kind {
            0 | 1 | 7 => 1,
            2 | 3 => 2,
            4..=6 => 4,
            8 | 10..=12 => 8,
            _ => unreachable!(),
        };
        if count > reader.remaining() / minimum {
            bail!("metadata array exceeds the file size");
        }
        macro_rules! read_values {
            ($method:ident) => {
                (0..count)
                    .map(|_| reader.$method())
                    .collect::<Result<Vec<_>>>()?
            };
        }
        Ok(match kind {
            0 => Self::U8(read_values!(u8)),
            1 => Self::I8(read_values!(i8)),
            2 => Self::U16(read_values!(u16)),
            3 => Self::I16(read_values!(i16)),
            4 => Self::U32(read_values!(u32)),
            5 => Self::I32(read_values!(i32)),
            6 => Self::F32(read_values!(f32)),
            7 => Self::Bool(read_values!(bool)),
            8 => {
                let values = (0..count)
                    .map(|_| reader.string("metadata array string"))
                    .collect::<Result<Vec<_>>>()?;
                Self::String(values)
            }
            10 => Self::U64(read_values!(u64)),
            11 => Self::I64(read_values!(i64)),
            12 => Self::F64(read_values!(f64)),
            _ => unreachable!(),
        })
    }

    fn type_name(&self) -> &'static str {
        match self {
            Self::U8(_) => "u8 array",
            Self::I8(_) => "i8 array",
            Self::U16(_) => "u16 array",
            Self::I16(_) => "i16 array",
            Self::U32(_) => "u32 array",
            Self::I32(_) => "i32 array",
            Self::F32(_) => "f32 array",
            Self::Bool(_) => "bool array",
            Self::String(_) => "string array",
            Self::U64(_) => "u64 array",
            Self::I64(_) => "i64 array",
            Self::F64(_) => "f64 array",
        }
    }
}

struct Reader<'a> {
    bytes: &'a [u8],
    position: usize,
}

impl<'a> Reader<'a> {
    fn new(bytes: &'a [u8]) -> Self {
        Self { bytes, position: 0 }
    }

    fn position(&self) -> usize {
        self.position
    }

    fn remaining(&self) -> usize {
        self.bytes.len() - self.position
    }

    fn take(&mut self, length: usize) -> Result<&'a [u8]> {
        let end = self
            .position
            .checked_add(length)
            .ok_or_else(|| anyhow!("GGUF offset overflow"))?;
        if end > self.bytes.len() {
            bail!("truncated GGUF at byte {}", self.position);
        }
        let value = &self.bytes[self.position..end];
        self.position = end;
        Ok(value)
    }

    fn array<const N: usize>(&mut self) -> Result<[u8; N]> {
        self.take(N)?
            .try_into()
            .map_err(|_| anyhow!("invalid scalar width"))
    }

    fn u8(&mut self) -> Result<u8> {
        Ok(self.take(1)?[0])
    }
    fn i8(&mut self) -> Result<i8> {
        Ok(self.u8()?.cast_signed())
    }
    fn u16(&mut self) -> Result<u16> {
        Ok(u16::from_le_bytes(self.array()?))
    }
    fn i16(&mut self) -> Result<i16> {
        Ok(i16::from_le_bytes(self.array()?))
    }
    fn u32(&mut self) -> Result<u32> {
        Ok(u32::from_le_bytes(self.array()?))
    }
    fn i32(&mut self) -> Result<i32> {
        Ok(i32::from_le_bytes(self.array()?))
    }
    fn f32(&mut self) -> Result<f32> {
        Ok(f32::from_le_bytes(self.array()?))
    }
    fn u64(&mut self) -> Result<u64> {
        Ok(u64::from_le_bytes(self.array()?))
    }
    fn i64(&mut self) -> Result<i64> {
        Ok(i64::from_le_bytes(self.array()?))
    }
    fn f64(&mut self) -> Result<f64> {
        Ok(f64::from_le_bytes(self.array()?))
    }

    fn bool(&mut self) -> Result<bool> {
        match self.u8()? {
            0 => Ok(false),
            1 => Ok(true),
            value => bail!("invalid GGUF boolean {value}"),
        }
    }

    fn usize(&mut self, description: &str) -> Result<usize> {
        let value = self.u64()?;
        usize::try_from(value).map_err(|_| anyhow!("{description} does not fit usize"))
    }

    fn count(&mut self, description: &str) -> Result<usize> {
        self.usize(description)
    }

    fn string(&mut self, description: &str) -> Result<String> {
        let length = self.usize(&format!("{description} length"))?;
        let bytes = self.take(length)?;
        Ok(str::from_utf8(bytes)
            .with_context(|| format!("{description} is not UTF-8"))?
            .to_owned())
    }
}

fn validate_name(name: &str, description: &str) -> Result<()> {
    if name.is_empty() || name.contains('\0') {
        bail!("invalid {description}");
    }
    Ok(())
}

fn validate_architecture(metadata: &HashMap<String, MetadataValue>) -> Result<()> {
    match metadata.get("general.architecture") {
        Some(MetadataValue::String(value)) if !value.is_empty() => Ok(()),
        Some(MetadataValue::String(_)) => bail!("general.architecture is empty"),
        Some(value) => bail!("general.architecture is {}", value.type_name()),
        None => bail!("missing GGUF metadata `general.architecture`"),
    }
}

fn metadata_alignment(metadata: &HashMap<String, MetadataValue>) -> Result<usize> {
    let alignment = match metadata.get("general.alignment") {
        Some(MetadataValue::U32(value)) => usize::try_from(*value)?,
        Some(value) => bail!("general.alignment is {}", value.type_name()),
        None => DEFAULT_ALIGNMENT,
    };
    if !alignment.is_power_of_two() {
        bail!("GGUF alignment must be a nonzero power of two");
    }
    Ok(alignment)
}

fn tensor_byte_len(name: &str, dimensions: &[usize], kind: TensorType) -> Result<usize> {
    let elements = dimensions.iter().try_fold(1_usize, |count, dimension| {
        count
            .checked_mul(*dimension)
            .ok_or_else(|| anyhow!("tensor `{name}` size overflow"))
    })?;
    match kind {
        TensorType::F32 => elements
            .checked_mul(size_of::<f32>())
            .ok_or_else(|| anyhow!("tensor `{name}` byte size overflow")),
        TensorType::Bf16 => elements
            .checked_mul(size_of::<u16>())
            .ok_or_else(|| anyhow!("tensor `{name}` byte size overflow")),
        TensorType::Q8_0 => {
            if !dimensions[0].is_multiple_of(Q8_BLOCK_VALUES) {
                bail!("Q8_0 tensor `{name}` row width is not divisible by 32");
            }
            (elements / Q8_BLOCK_VALUES)
                .checked_mul(Q8_BLOCK_BYTES)
                .ok_or_else(|| anyhow!("tensor `{name}` byte size overflow"))
        }
        TensorType::Mxfp4 => {
            if !dimensions[0].is_multiple_of(MXFP4_BLOCK_VALUES) {
                bail!("MXFP4 tensor `{name}` row width is not divisible by 32");
            }
            (elements / MXFP4_BLOCK_VALUES)
                .checked_mul(MXFP4_BLOCK_BYTES)
                .ok_or_else(|| anyhow!("tensor `{name}` byte size overflow"))
        }
    }
}

fn align_up(value: usize, alignment: usize) -> Result<usize> {
    value
        .checked_add(alignment - 1)
        .map(|value| value & !(alignment - 1))
        .ok_or_else(|| anyhow!("GGUF alignment overflow"))
}

fn validate_ranges(
    bytes: &[u8],
    data_start: usize,
    tensors: &HashMap<String, TensorInfo>,
) -> Result<()> {
    if data_start > bytes.len() {
        bail!("tensor data starts past end of file");
    }
    let mut ranges = Vec::with_capacity(tensors.len());
    for (name, info) in tensors {
        let start = data_start
            .checked_add(info.offset)
            .ok_or_else(|| anyhow!("tensor `{name}` offset overflow"))?;
        let end = start
            .checked_add(info.byte_len)
            .ok_or_else(|| anyhow!("tensor `{name}` range overflow"))?;
        if end > bytes.len() {
            bail!("tensor `{name}` exceeds file size");
        }
        ranges.push((start, end, name));
    }
    ranges.sort_unstable_by_key(|range| range.0);
    for pair in ranges.windows(2) {
        if pair[0].1 > pair[1].0 {
            bail!("tensor `{}` overlaps tensor `{}`", pair[0].2, pair[1].2);
        }
    }
    Ok(())
}

fn type_error<T>(key: &str, expected: &str, value: &MetadataValue) -> Result<T> {
    bail!(
        "GGUF metadata `{key}` is {}, expected {expected}",
        value.type_name()
    )
}