From 5efca0f41718486ffdab5c0d05a9983f99c2a683 Mon Sep 17 00:00:00 2001
From: Morgan Thomas <morgan.a.s.thomas@gmail.com>
Date: Fri, 5 Apr 2024 18:14:24 -0400
Subject: [PATCH] wip: mul64 chip

---
 alu_u32/src/lib.rs           |  1 +
 alu_u32/src/mul/mod.rs       |  3 +-
 alu_u32/src/mul64/columns.rs | 29 +++++++++++++++++++
 alu_u32/src/mul64/mod.rs     | 56 ++++++++++++++++++++++++++++++++++++
 alu_u32/src/mul64/stark.rs   | 22 ++++++++++++++
 5 files changed, 109 insertions(+), 2 deletions(-)
 create mode 100644 alu_u32/src/mul64/columns.rs
 create mode 100644 alu_u32/src/mul64/mod.rs
 create mode 100644 alu_u32/src/mul64/stark.rs

diff --git a/alu_u32/src/lib.rs b/alu_u32/src/lib.rs
index 778fa0db..9f6c84fd 100644
--- a/alu_u32/src/lib.rs
+++ b/alu_u32/src/lib.rs
@@ -8,5 +8,6 @@ pub mod com;
 pub mod div;
 pub mod lt;
 pub mod mul;
+pub mod mul64;
 pub mod shift;
 pub mod sub;
diff --git a/alu_u32/src/mul/mod.rs b/alu_u32/src/mul/mod.rs
index dba1f987..a6ba2486 100644
--- a/alu_u32/src/mul/mod.rs
+++ b/alu_u32/src/mul/mod.rs
@@ -5,7 +5,7 @@ use alloc::vec::Vec;
 use columns::{Mul32Cols, MUL_COL_MAP, NUM_MUL_COLS};
 use valida_bus::MachineWithGeneralBus;
 use valida_cpu::MachineWithCpuChip;
-use valida_machine::{instructions, Chip, Instruction, Interaction, Mulhs, Mulhu, Operands, Word};
+use valida_machine::{instructions, Chip, Instruction, Interaction, Mulhs, Mulhu, Operands, StarkConfig, Word};
 use valida_opcodes::{MUL32, MULHS32, MULHU32};
 use valida_range::MachineWithRangeChip;
 
@@ -13,7 +13,6 @@ use core::borrow::BorrowMut;
 use p3_air::VirtualPairCol;
 use p3_field::{AbstractField, Field, PrimeField};
 use p3_matrix::dense::RowMajorMatrix;
-use valida_machine::StarkConfig;
 
 pub mod columns;
 pub mod stark;
diff --git a/alu_u32/src/mul64/columns.rs b/alu_u32/src/mul64/columns.rs
new file mode 100644
index 00000000..71e9f070
--- /dev/null
+++ b/alu_u32/src/mul64/columns.rs
@@ -0,0 +1,29 @@
+extern crate alloc;
+
+use alloc::borrow::{Borrow, BorrowMut};
+use core::mem::{size_of, transmute};
+use valida_derive::AlignedBorrow;
+use valida_machine::Word;
+use valida_util::indices_arr;
+
+#[derive(Default)]
+pub struct Word64<T> {
+    pub most_significant: Word<T>,
+    pub least_significant: Word<T>,
+}
+
+#[derive(AlignedBorrow, Default)]
+pub struct Mul64Cols<T> {
+    pub input_1: Word<T>,
+    pub input_2: Word<T>,
+    pub output: Word64<T>,
+    // TODO
+}
+
+pub const NUM_MUL64_COLS: usize = size_of::<Mul64Cols<u8>>();
+pub const MUL64_COL_MAP: Mul64Cols<usize> = make_col_map();
+
+const fn make_col_map() -> Mul64Cols<usize> {
+    let indices_arr = indices_arr::<NUM_MUL64_COLS>();
+    unsafe { transmute::<[usize; NUM_MUL64_COLS], Mul64Cols<usize>>(indices_arr) }
+}
diff --git a/alu_u32/src/mul64/mod.rs b/alu_u32/src/mul64/mod.rs
new file mode 100644
index 00000000..63b2fddd
--- /dev/null
+++ b/alu_u32/src/mul64/mod.rs
@@ -0,0 +1,56 @@
+extern crate alloc;
+
+pub mod columns;
+pub mod stark;
+
+use alloc::borrow::Borrow;
+use alloc::vec::Vec;
+use columns::{Mul64Cols, MUL64_COL_MAP, NUM_MUL64_COLS, Word64};
+use valida_machine::{Chip, Interaction, Machine, Word, StarkConfig};
+
+use p3_field::{AbstractField, Field, PrimeField};
+use p3_matrix::dense::RowMajorMatrix;
+
+pub struct Operation {
+    pub inputs: (Word<u8>, Word<u8>),
+    pub output: Word64<u8>,
+}
+
+pub struct Mul64Chip {
+    pub operations: Vec<Operation>,
+}
+
+impl<M, SC> Chip<M, SC> for Mul64Chip
+where
+    M: Machine<SC::Val>,
+    SC: StarkConfig,
+{
+    fn generate_trace(&self, _machine: &M) -> RowMajorMatrix<SC::Val> {
+        todo!()
+    }
+
+    fn global_receives(&self, machine: &M) -> Vec<Interaction<SC::Val>> {
+        todo!()
+    }
+}
+
+impl Mul64Chip {
+    fn op_to_row<F>(&self, op: &Operation, cosl: &mut Mul64Cols<F>)
+    where
+        F: PrimeField,
+    {
+        todo!()
+    }
+
+    fn set_cols<F>(&self, a: &Word<u8>, b: &Word<u8>, c: &Word64<u8>, cols: &mut Mul64Cols<F>)
+    where
+        F: PrimeField,
+    {
+        todo!()
+    }
+}
+
+pub trait MachineWithMul64Chip<F: Field>: Machine<F> {
+    fn mul_u64(&self) -> &Mul64Chip;
+    fn mul_u64_mut(&mut self) -> &mut Mul64Chip;
+}
diff --git a/alu_u32/src/mul64/stark.rs b/alu_u32/src/mul64/stark.rs
new file mode 100644
index 00000000..31aaf4d3
--- /dev/null
+++ b/alu_u32/src/mul64/stark.rs
@@ -0,0 +1,22 @@
+use super::columns::Mul64Cols;
+use super::Mul64Chip;
+
+use crate::mul64::columns::NUM_MUL64_COLS;
+use p3_air::{Air, AirBuilder, BaseAir};
+use p3_field::{PrimeField};
+
+impl<F> BaseAir<F> for Mul64Chip {
+    fn width(&self) -> usize {
+        NUM_MUL64_COLS
+    }
+}
+
+impl<F, AB> Air<AB> for Mul64Chip
+where
+    F: PrimeField,
+    AB: AirBuilder<F = F>,
+{
+    fn eval(&self, _builder: &mut AB) {
+        todo!()
+    }
+}