1use crate::tree::CommittedMerkleTree;
 2use crate::FieldExt;
 3
 4#[derive(Clone)]
 5pub struct TensorCode<F>(pub Vec<Vec<F>>)
 6where
 7    F: FieldExt;
 8
 9impl<F: FieldExt> TensorCode<F> {
10    pub fn commit(&self, num_cols: usize, num_rows: usize) -> CommittedTensorCode<F> {
11        // Flatten the tensor codeword in column major order
12        let mut tensor_codeword = vec![];
13        for j in 0..(num_cols * 2) {
14            for i in 0..num_rows {
15                tensor_codeword.push(self.0[i][j])
16            }
17        }
18        // Merkle commit the codewords
19        let committed_tree = CommittedMerkleTree::from_leaves(tensor_codeword, num_cols * 2);
20
21        CommittedTensorCode {
22            committed_tree,
23            tensor_codeword: Self(self.0.clone()),
24        }
25    }
26}
27
28#[derive(Clone)]
29pub struct CommittedTensorCode<F: FieldExt> {
30    pub committed_tree: CommittedMerkleTree<F>,
31    pub tensor_codeword: TensorCode<F>,
32}
33
34impl<F: FieldExt> CommittedTensorCode<F> {
35    pub fn query_column(&self, column: usize) -> Vec<F> {
36        let num_rows = self.tensor_codeword.0.len();
37
38        let leaves =
39            self.committed_tree.leaves[column * num_rows..((column + 1) * num_rows)].to_vec();
40        leaves
41    }
42}