lansky/src/vm.rs

use std::cell::Cell;

use crate::{audio::{Channel, Sample}, error::{Lansky, LanskyError, LanskyStackReason}, device::Proto, rack::Rack};

pub(crate) struct Procedure {
    proto: Proto,
    blocks: Vec<Block>,
    analysis: Cell<Option<Analysis>>,
}

#[derive(Clone, Copy)]
struct Analysis {
    stack_depth: usize,

    proto_n_inputs: usize,
    proto_n_outputs: usize,

    n_reg_ls: usize,
}

pub struct Block {
    pub name: String,
    pub body: Vec<BodyInstruction>,
    pub exit: ExitInstruction
}

pub enum BodyInstruction {
    Mov(MovType, MovDst, MovSrc),
    Range(RangeDst, i32, i32),
    Inp(InpSrc),
    Dup(),
    Oup(MovType, OupDst),
    Pop(OupDst),
    New(MovSrc, MovSrc, MovSrc),
    Invoke(String, usize, usize),
}

pub enum ExitInstruction {
    Exit,
}

pub enum MovType { Mov, Add, Xor }

#[derive(Clone, Copy)] pub enum MovDst { L(RegL), R(RegR), Z(RegZ), Ix }
#[derive(Clone, Copy)] pub enum MovSrc { A(RegA), L(RegL), R(RegR), Z(RegZ), Ix, Imm(Sample) }
#[derive(Clone, Copy)] pub enum RangeDst { A(RegA), L(RegL), Z(RegZ) }
#[derive(Clone, Copy)] pub enum InpSrc { A(RegA), L(RegL), Z(RegZ) }
#[derive(Clone, Copy)] pub enum OupDst { L(RegL), Z(RegZ) }

impl OupDst {
    fn to_mov_dst(self) -> MovDst {
        match self {
            Self::L(l) => MovDst::L(l),
            Self::Z(z) => MovDst::Z(z),
        }
    }
}

#[derive(Clone, Copy)] pub struct RegA(pub u8);
#[derive(Clone, Copy)] pub struct RegL(pub u8);
#[derive(Clone, Copy)] pub struct RegR(pub u8);
#[derive(Clone, Copy)] pub struct RegZ(pub u8);

struct Scope {
    ix: Sample,

    reg_as: Vec<Channel>,
    reg_ls: Vec<Channel>,
    reg_rs: [Sample; 256],
    reg_zs: Vec<Channel>,

    loop_as: [LoopRange; 256],
    loop_ls: [LoopRange; 256],
    loop_zs: [LoopRange; 256],
}

#[derive(Clone, Copy)]
struct LoopRange {
    start: i32,
    length: i32, 
}

impl Default for LoopRange {
    fn default() -> Self {
        Self { start: 0, length: 0 }
    }
}

impl Scope {
    fn new(analysis: Analysis, inputs: Vec<Channel>) -> Scope {
        assert_eq!(inputs.len(), analysis.proto_n_inputs); // analysis is done elsewhere

        let reg_as: Vec<Channel> = inputs;
        let mut reg_ls: Vec<Channel> = vec![];
        let reg_rs: [Sample; 256] = [Sample::default(); 256];
        let mut reg_zs: Vec<Channel> = vec![];

        for _ in 0..analysis.n_reg_ls { reg_ls.push(Channel::blank()); }
        for _ in 0..analysis.proto_n_outputs { reg_zs.push(Channel::blank()); }

        Scope {
            ix: Sample(0),

            reg_as, 
            reg_ls, 
            reg_rs, 
            reg_zs,

            loop_as: [LoopRange::default(); 256],
            loop_ls: [LoopRange::default(); 256],
            loop_zs: [LoopRange::default(); 256],
        }
    }

    fn exit(self) -> Vec<Channel> {
        self.reg_zs
    }

    fn read(&self, src: MovSrc) -> Sample {
        let flat_ix = self.ix.0 as i32;
        let (chan, loop_) = match src {
            MovSrc::A(RegA(a)) => {
                let a = a as usize;
                (&self.reg_as[a], &self.loop_as[a])
            }
            MovSrc::L(RegL(l)) => {
                let l = l as usize;
                (&self.reg_ls[l], &self.loop_ls[l])
            }
            MovSrc::R(RegR(r)) => {
                let r = r as usize;
                return self.reg_rs[r]
            }
            MovSrc::Z(RegZ(z)) => {
                let z = z as usize;
                (&self.reg_zs[z], &self.loop_zs[z])
            }
            MovSrc::Ix => return self.ix,
            MovSrc::Imm(i) => return i,
        };

        let mut ix = flat_ix;
        if loop_.length != 0 { ix %= loop_.length; }
        ix += chan.formal().start + loop_.start;

        chan[ix]
    }

    fn write(&mut self, ty: &MovType, dst: MovDst, value: Sample, offset: i32) {
        let flat_ix = self.ix.0 as i32;
        let (chan, loop_) = match dst {
            MovDst::L(RegL(l)) => {
                let l = l as usize;
                (&mut self.reg_ls[l], &self.loop_ls[l])
            }
            MovDst::R(RegR(r)) => {
                let r = r as usize;
                self.reg_rs[r] = value; return
            }
            MovDst::Z(RegZ(z)) => {
                let z = z as usize;
                (&mut self.reg_zs[z], &self.loop_zs[z])
            }
            MovDst::Ix => { self.ix = value; return }
        };

        let mut ix = flat_ix + offset;
        if loop_.length != 0 { ix %= loop_.length; }
        ix += chan.formal().start + loop_.start;

        chan[ix] = value
    }

    fn replace(&mut self, dst: OupDst, value: Channel) {
        match dst {
            OupDst::L(RegL(l)) => { 
                self.reg_ls[l as usize] = value; 
                self.loop_ls[l as usize] = LoopRange::default() 
            }
            OupDst::Z(RegZ(z)) => { 
                self.reg_zs[z as usize] = value; 
                self.loop_zs[z as usize] = LoopRange::default() 
            }
        }
    }

    fn apply_range(&mut self, dst: RangeDst, start: i32, length: i32) {
        let lr = LoopRange{start, length};
        match dst {
            RangeDst::A(RegA(a)) => {
                self.loop_as[a as usize] = lr
            }
            RangeDst::L(RegL(l)) => {
                self.loop_ls[l as usize] = lr
            }
            RangeDst::Z(RegZ(z)) => {
                self.loop_zs[z as usize] = lr
            }
        }
    }

    fn inp(&self, src: InpSrc) -> Channel {
        let (chan, loop_) = match src {
            InpSrc::A(RegA(a)) => {
                let a = a as usize;
                (&self.reg_as[a], &self.loop_as[a])
            }
            InpSrc::L(RegL(l)) => {
                let l = l as usize;
                (&self.reg_ls[l], &self.loop_ls[l])
            }
            InpSrc::Z(RegZ(z)) => {
                let z = z as usize;
                (&self.reg_zs[z], &self.loop_zs[z])
            }
        };
        chan.slice(loop_.start, loop_.length)
    }
}

impl Procedure {
    pub(crate) fn new(proto: Proto, blocks: Vec<Block>) -> Procedure {
        Procedure {
            proto,
            blocks,
            analysis: Cell::new(None),
        }
    }

    pub fn invoke(&self, rack: &Rack, inputs: Vec<Channel>) -> Lansky<Vec<Channel>> {
        let analysis = self.internal_analyze(rack)?;

        let mut scope = Scope::new(analysis, inputs);
        let mut stack: Vec<Channel> = vec![];

        let block_ix = 0;
        loop {
            let block = &self.blocks[block_ix];
            for instruction in block.body.iter() {
                match instruction {
                    BodyInstruction::Mov(ty, dst, src) => {
                        scope.write(ty, *dst, scope.read(*src), 0)
                    }
                    BodyInstruction::Range(dst, start, len) => {
                        scope.apply_range(*dst, *start, *len)
                    }
                    BodyInstruction::Inp(src) => {
                        let slice = scope.inp(*src);
                        stack.push(slice)
                    }
                    BodyInstruction::Dup() => {
                        let top = stack.pop().expect("stack can't be empty");
                        stack.push(top.clone())
                    }
                    BodyInstruction::Oup(ty, dst) => {
                        let top = stack.pop().expect("stack can't be empty");
                        let bounds = top.bounds();
                        for ix in bounds.informal {
                            let offset = ix - bounds.formal.start;
                            scope.write(ty, dst.to_mov_dst(), top[ix], offset)
                        }
                    }
                    BodyInstruction::Pop(dst) => {
                        let top = stack.pop().expect("stack can't be empty");
                        scope.replace(*dst, top)
                    }
                    BodyInstruction::New(formal_start, formal_len, fill) => {
                        let start = scope.read(*formal_start);
                        let len = scope.read(*formal_len);
                        let fill = scope.read(*fill);
                        let mut chan = Channel::new((start.0 as i32)..((start.0+len.0) as i32));
                        for i in chan.formal() {
                            chan[i] = fill
                        }
                        stack.push(chan)
                    }
                    BodyInstruction::Invoke(device, n_inputs, n_outputs) => {
                        let mut inputs: Vec<Channel> = vec![]; 
                        std::mem::swap(&mut stack, &mut inputs);
                        assert_eq!(inputs.len(), *n_inputs);
                        let outputs = rack.invoke_positional(device, inputs)?;
                        assert_eq!(outputs.len(), *n_outputs);
                        stack.extend(outputs);
                        stack.reverse()
                    }
                }
            }

            match block.exit {
                ExitInstruction::Exit => return Ok(scope.exit())
            }
        }
    }

    pub(crate) fn analyze(&self, rack: &Rack) -> Lansky<()> {
        self.internal_analyze(rack)?;
        Ok(())
    }

    fn internal_analyze(&self, rack: &Rack) -> Lansky<Analysis> {
        if let Some(x) = self.analysis.get() { return Ok(x) }

        let mut analysis = Analysis { 
            stack_depth: 0, 

            proto_n_inputs: self.proto.inputs.len(),
            proto_n_outputs: self.proto.outputs.len(),

            n_reg_ls: 0 
        };

        if self.blocks.len() == 0 {
            return Err(LanskyError::NoFirstBlock);
        }

        for block in self.blocks.iter() {
            block.touch_analysis(
                &mut analysis, &|d| rack.get_proto(d).cloned()
            )?;
        }

        self.analysis.set(Some(analysis));

        return Ok(analysis)
    }
}


impl Block {
    fn touch_analysis(&self, analysis: &mut Analysis, get_proto: &impl Fn(&str) -> Lansky<Proto>) -> Lansky<()> {
        for instruction in self.body.iter() {
            instruction.touch_analysis(analysis, get_proto)?;
        }

        if analysis.stack_depth != 0 {
            return Err(LanskyError::WrongNStack { 
                why: LanskyStackReason::BlockMustEndWithEmptyStack,
                got: analysis.stack_depth, expected: 0 
            })
        }
        self.exit.touch_analysis(analysis)?;

        Ok(())
    }
}

impl BodyInstruction {
    fn touch_analysis(&self, analysis: &mut Analysis, get_proto: &impl Fn(&str) -> Lansky<Proto>) -> Lansky<()> {
        match self {
            BodyInstruction::Mov(_, dst, src) => {
                dst.touch_analysis(analysis)?;
                src.touch_analysis(analysis)?;
            }
            BodyInstruction::Range(dst, _, _) => {
                dst.touch_analysis(analysis)?;
            }
            BodyInstruction::Inp(src) => {
                analysis.stack_depth += 1;
                src.touch_analysis(analysis)?;
            }
            BodyInstruction::Dup() => {
                if analysis.stack_depth == 0 {
                    return Err(LanskyError::WrongNStack { 
                        why: LanskyStackReason::CannotDupEmptyStack,
                        expected: 1, got: 0 
                    })
                }
                analysis.stack_depth += 1;
            }
            BodyInstruction::Oup(_, dst) => {
                if analysis.stack_depth == 0 {
                    return Err(LanskyError::WrongNStack { 
                        why: LanskyStackReason::CannotOupEmptyStack,
                        expected: 1, got: 0 
                    })
                }
                analysis.stack_depth -= 1;

                dst.touch_analysis(analysis)?;
            }
            BodyInstruction::Pop(dst) => {
                if analysis.stack_depth == 0 {
                    return Err(LanskyError::WrongNStack { 
                        why: LanskyStackReason::CannotPopEmptyStack,
                        expected: 1, got: 0 
                    })
                }
                analysis.stack_depth -= 1;

                dst.touch_analysis(analysis)?;
            }
            BodyInstruction::New(formal_start, formal_len, fill) => {
                analysis.stack_depth += 1;

                formal_start.touch_analysis(analysis)?;
                formal_len.touch_analysis(analysis)?;
                fill.touch_analysis(analysis)?;
            }
            BodyInstruction::Invoke(device, n_inputs, n_outputs) => {
                let proto = get_proto(device)?;

                let actual_n_inputs = proto.inputs.len();
                let actual_n_outputs = proto.outputs.len();

                if *n_inputs != actual_n_inputs {
                    return Err(LanskyError::WrongNInputs { expected: actual_n_inputs, got: *n_inputs })
                }

                if *n_outputs != actual_n_outputs {
                    return Err(LanskyError::WrongNInputs { expected: actual_n_inputs, got: *n_outputs })
                }

                if analysis.stack_depth != *n_inputs {
                    return Err(LanskyError::WrongNStack {
                        why: LanskyStackReason::InvokeNInputs,
                        expected: *n_inputs,
                        got: analysis.stack_depth,
                    })
                }
                analysis.stack_depth = *n_outputs;
            }
        }
        Ok(())
    }
}

impl ExitInstruction {
    fn touch_analysis(&self, analysis: &mut Analysis) -> Lansky<()> {
        match self {
            ExitInstruction::Exit => {}
        }
        Ok(())
    }
}

impl MovDst {
    fn touch_analysis(&self, analysis: &mut Analysis) -> Lansky<()> {
        match self {
            Self::L(l) => l.touch_analysis(analysis),
            Self::R(r) => r.touch_analysis(analysis),
            Self::Z(z) => z.touch_analysis(analysis),
            Self::Ix => Ok(()),
        }
    }
}

impl MovSrc {
    fn touch_analysis(&self, analysis: &mut Analysis) -> Lansky<()> {
        match self {
            Self::A(a) => a.touch_analysis(analysis),
            Self::L(l) => l.touch_analysis(analysis),
            Self::R(r) => r.touch_analysis(analysis),
            Self::Z(z) => z.touch_analysis(analysis),
            Self::Ix => Ok(()),
            Self::Imm(_) => Ok(()),
        }
    }
}

impl RangeDst {
    fn touch_analysis(&self, analysis: &mut Analysis) -> Lansky<()> {
        match self {
            Self::A(a) => a.touch_analysis(analysis),
            Self::L(l) => l.touch_analysis(analysis),
            Self::Z(z) => z.touch_analysis(analysis),
        }
    }
}

impl InpSrc {
    fn touch_analysis(&self, analysis: &mut Analysis) -> Lansky<()> {
        match self {
            Self::A(a) => a.touch_analysis(analysis),
            Self::L(l) => l.touch_analysis(analysis),
            Self::Z(z) => z.touch_analysis(analysis),
        }
    }
}

impl OupDst {
    fn touch_analysis(&self, analysis: &mut Analysis) -> Lansky<()> {
        match self {
            Self::L(l) => l.touch_analysis(analysis),
            Self::Z(z) => z.touch_analysis(analysis),
        }
    }
}


impl RegA {
    fn touch_analysis(&self, analysis: &mut Analysis) -> Lansky<()> {
        if self.0 as usize >= analysis.proto_n_inputs {
            return Err(LanskyError::MissingRegA { i: self.0 as usize });
        }
        Ok(())
    }
}

impl RegL {
    fn touch_analysis(&self, analysis: &mut Analysis) -> Lansky<()> {
        analysis.n_reg_ls = analysis.n_reg_ls.max(self.0 as usize);
        Ok(())
    }
}

impl RegR {
    fn touch_analysis(&self, analysis: &mut Analysis) -> Lansky<()> {
        Ok(())
    }
}

impl RegZ {
    fn touch_analysis(&self, analysis: &mut Analysis) -> Lansky<()> {
        if self.0 as usize >= analysis.proto_n_outputs {
            return Err(LanskyError::MissingRegZ { i: self.0 as usize });
        }
        Ok(())
    }
}