polydat 0.2.0

// Copyright 2024-2026 Jonathan Shook
// SPDX-License-Identifier: Apache-2.0

//! JIT kernel types: structs and impls for all four kernel variants.
//!
//! `JitCore` holds the shared buffer, slot map, and module handle.
//! The four kernel structs (`JitKernelRaw`, `JitKernelPush`,
//! `JitKernelPull`, `JitKernelPushPull`) wrap a `JitCore` and a
//! compiled function pointer, providing `eval` and accessor methods.

use std::collections::HashMap;

use cranelift_jit::JITModule;

use crate::ast::PolydatNode;
use crate::kernel::ProvMask;

/// Shared fields for all JIT kernel variants.
pub(super) struct JitCore {
    pub(super) buffer: Vec<u64>,
    pub(super) coord_count: usize,
    pub(super) output_map: HashMap<String, usize>,
    pub(super) _module: JITModule,
    pub(super) _nodes: Vec<Box<dyn PolydatNode>>,
}

/// Compute slot provenance from input_dependents.
///
/// `input_dependents` is indexed by coordinate SLOT (callers
/// expand per-input lists across multi-slot inputs per §8.4
/// layer 1); `step_output_slots` carries every step's flattened
/// output slot list so multi-slot outputs share their step's
/// provenance word.
pub(super) fn compute_jit_slot_provenance(
    coord_count: usize,
    buffer_len: usize,
    step_output_slots: &[Vec<usize>],
    input_dependents: &[Vec<usize>],
) -> Vec<ProvMask> {
    let step_count = step_output_slots.len();
    let mut step_prov: Vec<ProvMask> =
        (0..step_count).map(|_| ProvMask::empty()).collect();
    for (input_slot, deps) in input_dependents.iter().enumerate() {
        for &step_idx in deps {
            if step_idx < step_count {
                step_prov[step_idx].set(input_slot);
            }
        }
    }
    let mut slot_prov: Vec<ProvMask> =
        (0..buffer_len).map(|_| ProvMask::empty()).collect();
    for (i, slot) in slot_prov.iter_mut().enumerate().take(coord_count) {
        slot.set(i);
    }
    for (i, outs) in step_output_slots.iter().enumerate() {
        for &slot in outs {
            if slot < slot_prov.len() {
                slot_prov[slot] = step_prov[i].clone();
            }
        }
    }
    slot_prov
}

macro_rules! jit_accessors {
    () => {
        /// Returns the number of coordinate inputs this kernel accepts.
        pub fn coord_count(&self) -> usize { self.core.coord_count }

        /// Returns the buffer slot index for the named output, if present.
        pub fn resolve_output(&self, name: &str) -> Option<usize> {
            self.core.output_map.get(name).copied()
        }

        /// Returns the raw u64 value stored in the named output slot.
        #[inline]
        pub fn get(&self, name: &str) -> u64 {
            self.core.buffer[self.core.output_map[name]]
        }

        /// Returns the raw u64 value stored at the given buffer slot index.
        #[inline]
        pub fn get_slot(&self, slot: usize) -> u64 {
            self.core.buffer[slot]
        }
    };
}

// ── JitKernelRaw ───────────────────────────────────────────

/// Raw JIT kernel: no provenance, all nodes evaluate unconditionally.
pub struct JitKernelRaw {
    pub(super) core: JitCore,
    pub(super) code_fn: unsafe fn(*const u64, *mut u64),
}

impl JitKernelRaw {
    /// Evaluate the kernel with the given coordinate values.
    ///
    /// Predicate violations (`is_positive`, `in_range`,
    /// `is_one_of`) from JIT-lowered code surface as normal
    /// Rust panics carrying the violation message. The
    /// longjmp wrapper in [`super::codegen::invoke_with_catch`]
    /// handles the transition back to Rust land.
    #[inline]
    pub fn eval(&mut self, coords: &[u64]) {
        self.core.buffer[..self.core.coord_count.min(coords.len())]
            .copy_from_slice(&coords[..self.core.coord_count.min(coords.len())]);
        let code_fn = self.code_fn;
        let buf_ptr_const = self.core.buffer.as_ptr();
        let buf_ptr_mut = self.core.buffer.as_mut_ptr();
        super::codegen::invoke_with_catch(move || {
            unsafe { (code_fn)(buf_ptr_const, buf_ptr_mut); }
        });
    }

    /// Evaluate and return the value at the given buffer slot index.
    #[inline]
    pub fn eval_for_slot(&mut self, coords: &[u64], slot: usize) -> u64 {
        self.eval(coords);
        self.core.buffer[slot]
    }

    /// Decompose into raw parts for hybrid kernel integration.
    pub fn into_parts(self) -> (unsafe fn(*const u64, *mut u64), JITModule) {
        (self.code_fn, self.core._module)
    }

    jit_accessors!();
}

// ── JitKernelPush ──────────────────────────────────────────

/// Push-only JIT kernel: per-node dirty tracking, no cone guard.
pub struct JitKernelPush {
    pub(super) core: JitCore,
    pub(super) code_fn_prov: unsafe fn(*const u64, *mut u64, *mut u8),
    pub(super) node_clean: Vec<u8>,
    pub(super) input_dependents: Vec<Vec<usize>>,
}

impl JitKernelPush {
    #[inline]
    fn set_inputs(&mut self, coords: &[u64]) {
        for (i, &c) in coords.iter().enumerate().take(self.core.coord_count) {
            if self.core.buffer[i] != c {
                self.core.buffer[i] = c;
                if i < self.input_dependents.len() {
                    for &step_idx in &self.input_dependents[i] {
                        self.node_clean[step_idx] = 0;
                    }
                }
            }
        }
    }

    /// Evaluate the kernel with the given coordinate values.
    #[inline]
    pub fn eval(&mut self, coords: &[u64]) {
        self.set_inputs(coords);
        let code_fn = self.code_fn_prov;
        let buf_const = self.core.buffer.as_ptr();
        let buf_mut = self.core.buffer.as_mut_ptr();
        let clean_mut = self.node_clean.as_mut_ptr();
        super::codegen::invoke_with_catch(move || {
            unsafe { (code_fn)(buf_const, buf_mut, clean_mut); }
        });
    }

    /// Evaluate and return the value at the given buffer slot index.
    #[inline]
    pub fn eval_for_slot(&mut self, coords: &[u64], slot: usize) -> u64 {
        self.eval(coords);
        self.core.buffer[slot]
    }

    jit_accessors!();
}

// ── JitKernelPull ──────────────────────────────────────────

/// Pull-only JIT kernel: cone guard, but all nodes run when cone is dirty.
/// Uses the raw (non-provenance) JIT function — no per-node clean checks.
pub struct JitKernelPull {
    pub(super) core: JitCore,
    pub(super) code_fn: unsafe fn(*const u64, *mut u64),
    pub(super) slot_provenance: Vec<ProvMask>,
    pub(super) changed_mask: ProvMask,
}

impl JitKernelPull {
    #[inline]
    fn set_inputs(&mut self, coords: &[u64]) {
        self.changed_mask.clear();
        for (i, &c) in coords.iter().enumerate().take(self.core.coord_count) {
            if self.core.buffer[i] != c {
                self.core.buffer[i] = c;
                self.changed_mask.set(i);
            }
        }
    }

    /// Evaluate the kernel with the given coordinate values.
    #[inline]
    pub fn eval(&mut self, coords: &[u64]) {
        self.set_inputs(coords);
        let code_fn = self.code_fn;
        let buf_const = self.core.buffer.as_ptr();
        let buf_mut = self.core.buffer.as_mut_ptr();
        super::codegen::invoke_with_catch(move || {
            unsafe { (code_fn)(buf_const, buf_mut); }
        });
    }

    /// Evaluate and return the value at the given buffer slot index,
    /// skipping evaluation if the slot's provenance cone is unaffected.
    #[inline]
    pub fn eval_for_slot(&mut self, coords: &[u64], slot: usize) -> u64 {
        self.set_inputs(coords);
        if slot < self.slot_provenance.len()
            && !self.slot_provenance[slot].intersects(&self.changed_mask) {
            return self.core.buffer[slot];
        }
        let code_fn = self.code_fn;
        let buf_const = self.core.buffer.as_ptr();
        let buf_mut = self.core.buffer.as_mut_ptr();
        super::codegen::invoke_with_catch(move || {
            unsafe { (code_fn)(buf_const, buf_mut); }
        });
        self.core.buffer[slot]
    }

    jit_accessors!();
}

// ── JitKernelPushPull ──────────────────────────────────────

/// Full optimization: push-side dirty tracking + pull-side cone guard.
pub struct JitKernelPushPull {
    pub(super) core: JitCore,
    pub(super) code_fn_prov: unsafe fn(*const u64, *mut u64, *mut u8),
    pub(super) node_clean: Vec<u8>,
    pub(super) input_dependents: Vec<Vec<usize>>,
    pub(super) slot_provenance: Vec<ProvMask>,
    pub(super) changed_mask: ProvMask,
}

impl JitKernelPushPull {
    #[inline]
    fn set_inputs(&mut self, coords: &[u64]) {
        self.changed_mask.clear();
        for (i, &c) in coords.iter().enumerate().take(self.core.coord_count) {
            if self.core.buffer[i] != c {
                self.core.buffer[i] = c;
                self.changed_mask.set(i);
                if i < self.input_dependents.len() {
                    for &step_idx in &self.input_dependents[i] {
                        self.node_clean[step_idx] = 0;
                    }
                }
            }
        }
    }

    /// Evaluate the kernel with the given coordinate values.
    #[inline]
    pub fn eval(&mut self, coords: &[u64]) {
        self.set_inputs(coords);
        let code_fn = self.code_fn_prov;
        let buf_const = self.core.buffer.as_ptr();
        let buf_mut = self.core.buffer.as_mut_ptr();
        let clean_mut = self.node_clean.as_mut_ptr();
        super::codegen::invoke_with_catch(move || {
            unsafe { (code_fn)(buf_const, buf_mut, clean_mut); }
        });
    }

    /// Evaluate and return the value at the given buffer slot index,
    /// applying both push and pull optimizations.
    #[inline]
    pub fn eval_for_slot(&mut self, coords: &[u64], slot: usize) -> u64 {
        self.set_inputs(coords);
        if slot < self.slot_provenance.len()
            && !self.slot_provenance[slot].intersects(&self.changed_mask) {
            return self.core.buffer[slot];
        }
        let code_fn = self.code_fn_prov;
        let buf_const = self.core.buffer.as_ptr();
        let buf_mut = self.core.buffer.as_mut_ptr();
        let clean_mut = self.node_clean.as_mut_ptr();
        super::codegen::invoke_with_catch(move || {
            unsafe { (code_fn)(buf_const, buf_mut, clean_mut); }
        });
        self.core.buffer[slot]
    }

    jit_accessors!();
}