qudit-tensor 0.3.4

Accelerated and Extensible Quantum Library
Documentation
use std::collections::HashMap;

use qudit_core::ParamInfo;
use qudit_expr::ExpressionId;

#[derive(Clone)]
pub enum BytecodeInstruction {
    Write(ExpressionId, ParamInfo, usize),
    Matmul(usize, usize, usize, ParamInfo, ParamInfo),
    Kron(usize, usize, usize, ParamInfo, ParamInfo),
    Hadamard(usize, usize, usize, ParamInfo, ParamInfo),
    FRPR(usize, Vec<usize>, Vec<usize>, usize),
    Trace(usize, Vec<(usize, usize)>, usize),
}

impl std::fmt::Debug for BytecodeInstruction {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            BytecodeInstruction::Write(id, _, index) => {
                write!(f, "Write {} {:?}", id, index)
            }
            BytecodeInstruction::Matmul(a, b, c, _, _) => {
                write!(f, "Matmul {:?} {:?} {:?}", a, b, c)
            }
            BytecodeInstruction::Kron(a, b, c, _, _) => {
                write!(f, "Kron {:?} {:?} {:?}", a, b, c)
            }
            BytecodeInstruction::Hadamard(a, b, c, _, _) => {
                write!(f, "Kron {:?} {:?} {:?}", a, b, c)
            }
            BytecodeInstruction::FRPR(a, _, _, d) => {
                write!(f, "FRPR {:?} {:?}", a, d)
            }
            BytecodeInstruction::Trace(a, _, c) => {
                write!(f, "Trace {:?} {:?}", a, c)
            }
        }
    }
}

impl BytecodeInstruction {
    pub fn offset_buffer_indices(&mut self, offset: usize) {
        match self {
            BytecodeInstruction::Write(_, _, index) => {
                *index += offset;
            }
            BytecodeInstruction::Matmul(a, b, c, _, _) => {
                *a += offset;
                *b += offset;
                *c += offset;
            }
            BytecodeInstruction::Kron(a, b, c, _, _) => {
                *a += offset;
                *b += offset;
                *c += offset;
            }
            BytecodeInstruction::Hadamard(a, b, c, _, _) => {
                *a += offset;
                *b += offset;
                *c += offset;
            }
            BytecodeInstruction::FRPR(a, _, _, d) => {
                *a += offset;
                *d += offset;
            }
            BytecodeInstruction::Trace(a, _, c) => {
                *a += offset;
                *c += offset;
            }
        }
    }

    pub fn replace_buffer_indices(&mut self, buffer_map: &HashMap<usize, usize>) {
        match self {
            BytecodeInstruction::Write(_, _, index) => {
                if let Some(new_index) = buffer_map.get(index) {
                    *index = *new_index;
                }
            }
            BytecodeInstruction::Matmul(a, b, c, _, _) => {
                if let Some(new_index) = buffer_map.get(a) {
                    *a = *new_index;
                }
                if let Some(new_index) = buffer_map.get(b) {
                    *b = *new_index;
                }
                if let Some(new_index) = buffer_map.get(c) {
                    *c = *new_index;
                }
            }
            BytecodeInstruction::Kron(a, b, c, _, _) => {
                if let Some(new_index) = buffer_map.get(a) {
                    *a = *new_index;
                }
                if let Some(new_index) = buffer_map.get(b) {
                    *b = *new_index;
                }
                if let Some(new_index) = buffer_map.get(c) {
                    *c = *new_index;
                }
            }
            BytecodeInstruction::Hadamard(a, b, c, _, _) => {
                if let Some(new_index) = buffer_map.get(a) {
                    *a = *new_index;
                }
                if let Some(new_index) = buffer_map.get(b) {
                    *b = *new_index;
                }
                if let Some(new_index) = buffer_map.get(c) {
                    *c = *new_index;
                }
            }
            BytecodeInstruction::FRPR(a, _, _, d) => {
                if let Some(new_index) = buffer_map.get(a) {
                    *a = *new_index;
                }
                if let Some(new_index) = buffer_map.get(d) {
                    *d = *new_index;
                }
            }
            BytecodeInstruction::Trace(a, _, c) => {
                if let Some(new_index) = buffer_map.get(a) {
                    *a = *new_index;
                }
                if let Some(new_index) = buffer_map.get(c) {
                    *c = *new_index;
                }
            }
        }
    }
}