knok-build 0.4.0

use knok_core::{CallOp, Expr};

use super::{
    context::{tuple_id, TraceVars},
    expr::{call_output, unary},
    tensor::TraceTensor,
};

/// Records a reshape to the requested output tensor type.
pub fn reshape<Output: TraceTensor>(value: impl TraceTensor) -> Output {
    call_output(
        CallOp::Reshape(Output::tensor_type()),
        vec![value.into_expr()],
    )
}

/// Records a broadcast to the requested output tensor type.
pub fn broadcast<Output: TraceTensor>(value: impl TraceTensor) -> Output {
    call_output(
        CallOp::Broadcast(Output::tensor_type()),
        vec![value.into_expr()],
    )
}

/// Records a squeeze to the requested output tensor type.
pub fn squeeze<Output: TraceTensor>(value: impl TraceTensor) -> Output {
    call_output(
        CallOp::Squeeze(Output::tensor_type()),
        vec![value.into_expr()],
    )
}

/// Records an unsqueeze to the requested output tensor type.
pub fn unsqueeze<Output: TraceTensor>(value: impl TraceTensor) -> Output {
    call_output(
        CallOp::Unsqueeze(Output::tensor_type()),
        vec![value.into_expr()],
    )
}

/// Records a static slice from `starts` to the requested output tensor type.
pub fn slice<Output: TraceTensor, const N: usize>(
    value: impl TraceTensor,
    starts: [usize; N],
) -> Output {
    trace_slice(value, &starts)
}

fn trace_slice<Output: TraceTensor>(value: impl TraceTensor, starts: &[usize]) -> Output {
    call_output(
        CallOp::Slice {
            target: Output::tensor_type(),
            starts: starts.to_vec(),
        },
        vec![value.into_expr()],
    )
}

fn trace_pad<Output: TraceTensor>(value: impl TraceTensor, lows: &[usize]) -> Output {
    call_output(
        CallOp::Pad {
            target: Output::tensor_type(),
            lows: lows.to_vec(),
        },
        vec![value.into_expr()],
    )
}

/// Records low-side padding to the requested output tensor type.
pub fn pad<Output: TraceTensor, const N: usize>(
    value: impl TraceTensor,
    lows: [usize; N],
) -> Output {
    trace_pad(value, &lows)
}

/// Records a gather along `axis` to the requested output tensor type.
pub fn gather<Output: TraceTensor>(
    value: impl TraceTensor,
    indices: impl TraceTensor,
    axis: usize,
) -> Output {
    trace_gather(value, indices, axis)
}

fn trace_gather<Output: TraceTensor>(
    value: impl TraceTensor,
    indices: impl TraceTensor,
    axis: usize,
) -> Output {
    call_output(
        CallOp::Gather {
            target: Output::tensor_type(),
            axis,
        },
        vec![value.into_expr(), indices.into_expr()],
    )
}

/// Records a static take of one `index` along `axis`.
pub fn take<Output: TraceTensor>(value: impl TraceTensor, axis: usize, index: usize) -> Output {
    trace_take(value, axis, index)
}

fn trace_take<Output: TraceTensor>(value: impl TraceTensor, axis: usize, index: usize) -> Output {
    call_output(CallOp::Take { axis, index }, vec![value.into_expr()])
}

/// Records a take-along-axis using per-element indices.
pub fn take_along_axis<Output: TraceTensor>(
    value: impl TraceTensor,
    indices: impl TraceTensor,
    axis: usize,
) -> Output {
    trace_take_along_axis(value, indices, axis)
}

fn trace_take_along_axis<Output: TraceTensor>(
    value: impl TraceTensor,
    indices: impl TraceTensor,
    axis: usize,
) -> Output {
    call_output(
        CallOp::TakeAlongAxis { axis },
        vec![value.into_expr(), indices.into_expr()],
    )
}

/// Records concatenation of two tensors along `axis`.
pub fn concat<Output: TraceTensor>(
    lhs: impl TraceTensor,
    rhs: impl TraceTensor,
    axis: usize,
) -> Output {
    trace_concat(lhs, rhs, axis)
}

fn trace_concat<Output: TraceTensor>(
    lhs: impl TraceTensor,
    rhs: impl TraceTensor,
    axis: usize,
) -> Output {
    call_output(CallOp::Concat(axis), vec![lhs.into_expr(), rhs.into_expr()])
}

/// Records stacking of two tensors along `axis`.
pub fn stack<Output: TraceTensor>(
    lhs: impl TraceTensor,
    rhs: impl TraceTensor,
    axis: usize,
) -> Output {
    trace_stack(lhs, rhs, axis)
}

fn trace_stack<Output: TraceTensor>(
    lhs: impl TraceTensor,
    rhs: impl TraceTensor,
    axis: usize,
) -> Output {
    call_output(CallOp::Stack(axis), vec![lhs.into_expr(), rhs.into_expr()])
}

/// Records a split into statically sized sections along `axis`.
pub fn split<Output: TraceVars, const N: usize>(
    value: impl TraceTensor,
    axis: usize,
    sections: [usize; N],
) -> Output {
    trace_split(value, axis, &sections)
}

fn trace_split<Output: TraceVars>(
    value: impl TraceTensor,
    axis: usize,
    sections: &[usize],
) -> Output {
    if sections.len() != Output::var_count() {
        panic!(
            "split sections produce {} outputs, but the requested Rust output type has {} values",
            sections.len(),
            Output::var_count()
        );
    }
    let value = Expr::Call {
        op: CallOp::Split {
            axis,
            sections: sections.to_vec(),
        },
        args: vec![value.into_expr()],
    };
    Output::from_tuple_expr(tuple_id(), value)
}

/// Records tiling by static per-axis multiples.
pub fn tile<Output: TraceTensor, const N: usize>(
    value: impl TraceTensor,
    multiples: [usize; N],
) -> Output {
    trace_tile(value, &multiples)
}

fn trace_tile<Output: TraceTensor>(value: impl TraceTensor, multiples: &[usize]) -> Output {
    call_output(CallOp::Tile(multiples.to_vec()), vec![value.into_expr()])
}

/// Records repeating each element `count` times along `axis`.
pub fn repeat<Output: TraceTensor>(value: impl TraceTensor, axis: usize, count: usize) -> Output {
    trace_repeat(value, axis, count)
}

fn trace_repeat<Output: TraceTensor>(value: impl TraceTensor, axis: usize, count: usize) -> Output {
    call_output(CallOp::Repeat { axis, count }, vec![value.into_expr()])
}

/// Records reversal along every axis.
pub fn flip<T: TraceTensor>(value: T) -> T {
    trace_flip(value)
}

fn trace_flip<T: TraceTensor>(value: T) -> T {
    unary(CallOp::Flip(Vec::new()), value)
}

/// Records reversal along selected axes.
pub fn flip_axes<T: TraceTensor, const N: usize>(value: T, axes: [usize; N]) -> T {
    trace_flip_axes(value, &axes)
}

fn trace_flip_axes<T: TraceTensor>(value: T, axes: &[usize]) -> T {
    unary(CallOp::Flip(axes.to_vec()), value)
}

/// Records a static roll along one axis.
pub fn roll<T: TraceTensor>(value: T, axis: usize, shift: usize) -> T {
    trace_roll(value, axis, shift)
}

fn trace_roll<T: TraceTensor>(value: T, axis: usize, shift: usize) -> T {
    unary(CallOp::Roll { axis, shift }, value)
}

/// Records a default transpose to the requested output tensor type.
pub fn transpose<Output: TraceTensor>(value: impl TraceTensor) -> Output {
    trace_transpose(value)
}

fn trace_transpose<Output: TraceTensor>(value: impl TraceTensor) -> Output {
    call_output(CallOp::Transpose(Vec::new()), vec![value.into_expr()])
}

/// Records a transpose with explicit axes.
pub fn transpose_axes<Output: TraceTensor, const N: usize>(
    value: impl TraceTensor,
    axes: [usize; N],
) -> Output {
    trace_transpose_axes(value, &axes)
}

fn trace_transpose_axes<Output: TraceTensor>(value: impl TraceTensor, axes: &[usize]) -> Output {
    call_output(CallOp::Transpose(axes.to_vec()), vec![value.into_expr()])
}

/// Records a permutation with explicit axes and output type.
pub fn permute<Output: TraceTensor, const N: usize>(
    value: impl TraceTensor,
    axes: [usize; N],
) -> Output {
    trace_permute(value, &axes)
}

fn trace_permute<Output: TraceTensor>(value: impl TraceTensor, axes: &[usize]) -> Output {
    call_output(
        CallOp::Permute {
            target: Output::tensor_type(),
            axes: axes.to_vec(),
        },
        vec![value.into_expr()],
    )
}

/// Records a dimension permutation with output type inferred by type checking.
pub fn permute_dims<Output: TraceTensor, const N: usize>(
    value: impl TraceTensor,
    axes: [usize; N],
) -> Output {
    trace_permute_dims(value, &axes)
}

fn trace_permute_dims<Output: TraceTensor>(value: impl TraceTensor, axes: &[usize]) -> Output {
    call_output(CallOp::PermuteDims(axes.to_vec()), vec![value.into_expr()])
}

/// Records swapping two axes.
pub fn swapaxes<Output: TraceTensor>(
    value: impl TraceTensor,
    axis0: usize,
    axis1: usize,
) -> Output {
    trace_swapaxes(value, axis0, axis1)
}

fn trace_swapaxes<Output: TraceTensor>(
    value: impl TraceTensor,
    axis0: usize,
    axis1: usize,
) -> Output {
    call_output(CallOp::SwapAxes { axis0, axis1 }, vec![value.into_expr()])
}

/// Records moving one source axis to a destination position.
pub fn moveaxis<Output: TraceTensor>(
    value: impl TraceTensor,
    source: usize,
    destination: usize,
) -> Output {
    trace_moveaxis(value, source, destination)
}

fn trace_moveaxis<Output: TraceTensor>(
    value: impl TraceTensor,
    source: usize,
    destination: usize,
) -> Output {
    call_output(
        CallOp::MoveAxis {
            source,
            destination,
        },
        vec![value.into_expr()],
    )
}