wam 0.1.0

use std::collections::HashMap;

#[derive(Clone, PartialEq)]
enum Cell {
    Variable(Address),          // points to bounded term (or itself if unbound)
    Structure(Address),         // usize points to functor
    Functor(String, usize),     // String refers to the name, usize is the arity
}

impl std::fmt::Debug for Cell {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            &Cell::Variable(bind) => write!(f, "REF {:?}", bind),
            &Cell::Structure(ptr) => write!(f, "STR {:?}", ptr),
            &Cell::Functor(ref name, arity) => write!(f, "{}/{}", name, arity),
        }
    }
}

#[derive(Clone, Debug)]
pub enum Instruction {
    // query term instructions
    PutStructure(String, usize, Address),   // name, arity, register
    SetVariable(Address),                   // register
    SetValue(Address),                      // register
    
    // program instructions
    GetStructure(String, usize, Address),   // name, arity, register
    UnifyVariable(Address),                 // register
    UnifyValue(Address),                    // register

    // varaible argument instructions
    PutVariable(Address, Address),          // register, register
    PutValue(Address, Address),             // register, register
    GetVariable(Address, Address),          // register, register
    GetValue(Address, Address),             // register, register

    // controll instructions
    Call(String),                           // label
    Proceed,
}

enum Mode {
    Read,
    Write,
}

#[derive(Clone, PartialEq, Debug, Copy)]
pub enum Address {
    Heap(usize),
    Register(usize),
}

impl std::ops::Add<usize> for Address {
    type Output = Address;

    fn add(self, rhs: usize) -> Address {
        match self {
            Address::Heap(a) => Address::Heap(a + rhs),
            Address::Register(a) => Address::Register(a + rhs),
        }
    }
}

pub struct Instance {
    heap: Vec<Cell>,
    registers: Vec<Option<Cell>>,
    mode: Mode,
    s_ptr: Address,
    instructions: Vec<Instruction>,
    pc: usize,
    labels: HashMap<String, usize>,
}

const REGISTER_COUNT: usize = 100;

impl Instance {

    pub fn new() -> Instance {
        Instance {
            heap: Vec::new(),
            registers: vec![None; REGISTER_COUNT],
            mode: Mode::Read,
            s_ptr: Address::Heap(0),
            instructions: Vec::new(),
            pc: 0,
            labels: HashMap::new(),
        }
    }

    fn set_cell(&mut self, address: Address, cell: Cell) {
        match address {
            Address::Heap(a) => self.heap[a] = cell,
            Address::Register(a) => self.registers[a] = Some(cell),
        }
    }

    fn fetch_cell(&self, address: Address) -> Cell {
        match address {
            Address::Heap(a) => self.heap[a].clone(),
            Address::Register(a) => self.registers[a].clone().unwrap(),
        }
    }

    fn deref(&self, address: Address) -> Address {
        match self.fetch_cell(address) {
            Cell::Variable(v_address) => {
                if address == v_address {
                    address
                } else {
                    self.deref(v_address)
                }
            },
            _ => address,
        }
    }

    fn bind(&mut self, a1: Address, a2: Address) {
        let cells = (self.fetch_cell(a1), self.fetch_cell(a2));
        let (var_address, target_adresss) = match cells {
            (Cell::Variable(_), _) => (a1, a2),
            (_, Cell::Variable(_)) => (a2, a1),
            _ => panic!("neiter {:?} or {:?} points to a variable", a1, a2)
        };

        // Point the variable to the target cell
        self.set_cell(var_address, Cell::Variable(target_adresss));    
    }

    pub fn load(&mut self, instructions: Vec<Instruction>) {
        self.instructions.extend(instructions);
    }

    pub fn load_labeled(&mut self, label: String, instructions: Vec<Instruction>) {
        let index = self.instructions.len();
        self.load(instructions);
        self.labels.insert(label, index - 1);
    }

    fn inst_put_structure(&mut self, name: String, arity: usize, register: Address) {
        let str_index = self.heap.len();
        let str_cell = Cell::Structure(Address::Heap(str_index + 1));
        self.heap.push(str_cell.clone());
        self.heap.push(Cell::Functor(name.clone(), arity.clone()));               
        self.set_cell(register, str_cell);
    }

    fn inst_set_variable(&mut self, register: Address) {
        let ref_index = self.heap.len();
        let ref_cell = Cell::Variable(Address::Heap(ref_index));
        self.heap.push(ref_cell.clone());
        self.set_cell(register, ref_cell);
    }

    fn inst_set_value(&mut self, register: Address) {
        let cell = self.fetch_cell(register);
        self.heap.push(cell);
    }

    fn inst_get_structure(&mut self, name: String, arity: usize, register: Address) {
        let address = self.deref(register);
        match self.fetch_cell(address) {
            Cell::Variable(_) => {
                let str_index = self.heap.len();
                self.heap.push(Cell::Structure(Address::Heap(str_index + 1)));
                self.heap.push(Cell::Functor(name, arity));
                self.bind(address, Address::Heap(str_index));
                self.mode = Mode::Write;
            },
            Cell::Structure(a) => {
                let cell = self.fetch_cell(a);
                if cell == Cell::Functor(name, arity) {
                    self.s_ptr = a + 1;
                    self.mode = Mode::Read;
                } else {
                    panic!("fail");
                }
            },
            _ => panic!("fail"),
        }
    }

    fn inst_unify_variable(&mut self, register: Address) {
        match self.mode {
            Mode::Read => {
                let cell = self.fetch_cell(self.s_ptr);
                self.set_cell(register, cell);
            },
            Mode::Write => {
                let ref_index = self.heap.len();
                let cell = Cell::Variable(Address::Heap(ref_index));
                self.heap.push(cell.clone());
                self.set_cell(register, cell);
            }
        }
        
        self.s_ptr = self.s_ptr + 1;
    }

    fn inst_unift_value(&mut self, register: Address) {
        match self.mode {
            Mode::Read => {
                let s = self.s_ptr;
                self.unify(register, s);
            },
            Mode::Write => {
                let cell = self.fetch_cell(register);
                self.heap.push(cell);
            }
        }
    }

    fn inst_put_variable(&mut self, register_n: Address, register_i: Address) {
        let ref_index = self.heap.len();
        let cell = Cell::Variable(Address::Heap(ref_index));
        self.heap.push(cell.clone());
        self.set_cell(register_n, cell.clone());
        self.set_cell(register_i, cell);
    }

    fn inst_put_value(&mut self, register_n: Address, register_i: Address) {
        let cell = self.fetch_cell(register_n);
        self.set_cell(register_i, cell);
    }

    fn inst_get_variable(&mut self, register_n: Address, register_i: Address) {
        let cell = self.fetch_cell(register_i);
        self.set_cell(register_n, cell);
    }

    fn inst_get_value(&mut self, register_n: Address, register_i: Address) {
        self.unify(register_n, register_i);
    }

    fn inst_call(&mut self, label: String) {
        if self.labels.contains_key(&label) {
            self.pc = *self.labels.get(&label).unwrap();
        } else {
            panic!("Unification error");
        }
    }

    fn fetch_instruction(&self) -> Option<Instruction> {
        if self.instructions.len() <= self.pc {
            None
        } else {
            Some(self.instructions[self.pc].clone())
        }
    }

    pub fn execute(&mut self) {
        println!("Executing {:?}", self.fetch_instruction());

        match self.fetch_instruction() {
            Some(instruction) => match instruction {
                Instruction::PutStructure(name, arity, register) => self.inst_put_structure(name, arity, register),
                Instruction::SetVariable(register) => self.inst_set_variable(register),
                Instruction::SetValue(register) => self.inst_set_value(register),
                Instruction::GetStructure(name, arity, register) => self.inst_get_structure(name, arity, register),
                Instruction::UnifyVariable(register) => self.inst_unify_variable(register),
                Instruction::UnifyValue(register) => self.inst_unift_value(register),
                Instruction::PutVariable(register_n, register_i) => self.inst_put_variable(register_n, register_i),
                Instruction::PutValue(register_n, register_i) => self.inst_put_value(register_n, register_i),
                Instruction::GetVariable(register_n, register_i) => self.inst_get_variable(register_n, register_i),
                Instruction::GetValue(register_n, register_i) => self.inst_get_value(register_n, register_i),
                Instruction::Call(label) => self.inst_call(label),
                Instruction::Proceed => return,
            }
            None => return
        }

        self.pc += 1;
        self.execute();
    }

    fn unify(&mut self, ad1: Address, ad2: Address) {
        let mut pdl = Vec::new();
        pdl.push(ad1);
        pdl.push(ad2);

        while !pdl.is_empty() {
            let d1 = self.deref(pdl.pop().unwrap());
            let d2 = self.deref(pdl.pop().unwrap());
            println!("{:?} <> {:?}", d1, d2);
            if d1 != d2 {
                match (self.fetch_cell(d1), self.fetch_cell(d2)) {
                    (Cell::Variable(_), _) => self.bind(d1, d2),
                    (_, Cell::Variable(_)) => self.bind(d1, d2),
                    (Cell::Structure(str_1), Cell::Structure(str_2)) => {
                        match (self.fetch_cell(str_1), self.fetch_cell(str_2)) {
                            (Cell::Functor(f1, n1), Cell::Functor(f2, n2)) => {
                                if f1 != f2 || n1 != n2 {
                                    panic!("functors {}/{} and {}/{} do not match", f1, n1, f2, n2);
                                }
                                for i in 1..(n1+1) {
                                    pdl.push(str_1 + i);
                                    pdl.push(str_2 + i);
                                }
                            },
                            _ => panic!("structure cell addresses unwrapped to {:?} and {:?}", self.fetch_cell(str_1), self.fetch_cell(str_2)),
                        }
                    }
                    _ => {
                        panic!("cannot unify {:?} and {:?}", self.fetch_cell(d1), self.fetch_cell(d2));
                    },
                }
            }
        }
    }

    pub fn dump_heap(&self) {
        println!("HEAP:");
        for (index, cell) in self.heap.iter().enumerate() {
            println!("{:4} | {:?}", index, cell);
        }
    }

    pub fn dump_registers(&self) {
        println!("REGISTERS:");
        for (index, content) in self.registers.iter().enumerate() {
            println!("X_{} = {:?}", index, content);
        }
    }

}

#[cfg(test)]
mod tests {
    use super::*;

    fn instruction_set_1() -> Vec<Instruction> {
        vec![
            Instruction::PutStructure(String::from("h"), 2, Address::Register(3)),
            Instruction::SetVariable(Address::Register(2)),
            Instruction::SetVariable(Address::Register(5)),
            Instruction::PutStructure(String::from("f"), 1, Address::Register(4)),
            Instruction::SetValue(Address::Register(5)),
            Instruction::PutStructure(String::from("p"), 3, Address::Register(1)),
            Instruction::SetValue(Address::Register(2)),
            Instruction::SetValue(Address::Register(3)),
            Instruction::SetValue(Address::Register(4)),
        ]
    }

    fn instruction_set_2() -> Vec<Instruction> {
        vec![
            Instruction::GetStructure(String::from("p"), 3, Address::Register(1)),
            Instruction::UnifyVariable(Address::Register(2)),
            Instruction::UnifyVariable(Address::Register(3)),
            Instruction::UnifyVariable(Address::Register(4)),
            Instruction::GetStructure(String::from("f"), 1, Address::Register(2)),
            Instruction::UnifyVariable(Address::Register(5)),
            Instruction::GetStructure(String::from("h"), 2, Address::Register(3)),
            Instruction::UnifyValue(Address::Register(4)),
            Instruction::UnifyVariable(Address::Register(6)),
            Instruction::GetStructure(String::from("f"), 1, Address::Register(6)),
            Instruction::UnifyVariable(Address::Register(7)),
            Instruction::GetStructure(String::from("a"), 0, Address::Register(7)),
        ]
    }

    fn instruction_set_3() -> Vec<Instruction> {
        vec![
            Instruction::PutVariable(Address::Register(4), Address::Register(1)),
            Instruction::PutStructure(String::from("h"), 2, Address::Register(2)),
            Instruction::SetValue(Address::Register(4)),
            Instruction::SetVariable(Address::Register(5)),
            Instruction::PutStructure(String::from("f"), 1, Address::Register(3)),
            Instruction::SetValue(Address::Register(5)),
            Instruction::Call(String::from("p/3")),
        ]
    }

    fn instruction_set_4() -> (String, Vec<Instruction>) {
        (String::from("p/3"), vec![
            Instruction::GetStructure(String::from("f"), 1, Address::Register(1)),
            Instruction::Proceed,
            Instruction::UnifyVariable(Address::Register(1)),
            Instruction::GetStructure(String::from("h"), 2, Address::Register(2)),
            Instruction::UnifyVariable(Address::Register(5)),
            Instruction::UnifyVariable(Address::Register(6)),
            Instruction::GetValue(Address::Register(5), Address::Register(3)),
            Instruction::GetStructure(String::from("f"), 1, Address::Register(6)),
            Instruction::UnifyVariable(Address::Register(7)),
            Instruction::GetStructure(String::from("a"), 0, Address::Register(7)),
        ])
    }

    #[test]
    fn query_term_instructions() {
        let mut inst = Instance::new();
        inst.load(instruction_set_1());
        inst.execute();

        let expected_heap = vec![
            Cell::Structure(Address::Heap(1)),
            Cell::Functor(String::from("h"), 2),
            Cell::Variable(Address::Heap(2)),
            Cell::Variable(Address::Heap(3)),
            Cell::Structure(Address::Heap(5)),
            Cell::Functor(String::from("f"), 1),
            Cell::Variable(Address::Heap(3)),
            Cell::Structure(Address::Heap(8)),
            Cell::Functor(String::from("p"), 3),
            Cell::Variable(Address::Heap(2)),
            Cell::Structure(Address::Heap(1)),
            Cell::Structure(Address::Heap(5)),
        ];

        assert_eq!(inst.heap, expected_heap);
    }

    #[test]
    fn unify() {
        let instructions = vec![
            // f(X, g(X, a))
            // X1 = f(X2, X3)
            // X2 = X
            // X3 = g(X2, X4)
            // X4 = a
            Instruction::PutStructure(String::from("a"), 0, Address::Register(4)),
            Instruction::PutStructure(String::from("g"), 2, Address::Register(3)),
            Instruction::SetVariable(Address::Register(2)),
            Instruction::SetValue(Address::Register(4)),
            Instruction::PutStructure(String::from("f"), 2, Address::Register(1)),
            Instruction::SetValue(Address::Register(2)),
            Instruction::SetValue(Address::Register(3)),
            
            // f(b, Y)
            // X5 = f(X7, X8)
            // X6 = b
            // X7 = Y
            Instruction::PutStructure(String::from("b"), 0, Address::Register(6)),
            Instruction::PutStructure(String::from("f"), 2, Address::Register(5)),
            Instruction::SetValue(Address::Register(6)),
            Instruction::SetVariable(Address::Register(7)),
        ];

        let mut inst = Instance::new();
        inst.load(instructions);
        inst.execute();
        
        let expected_heap = vec![
            Cell::Structure(Address::Heap(1)),
            Cell::Functor(String::from("a"), 0),
            Cell::Structure(Address::Heap(3)),
            Cell::Functor(String::from("g"), 2),
            Cell::Variable(Address::Heap(4)),
            Cell::Structure(Address::Heap(1)),
            Cell::Structure(Address::Heap(7)),
            Cell::Functor(String::from("f"), 2),
            Cell::Variable(Address::Heap(4)),
            Cell::Structure(Address::Heap(3)),
            Cell::Structure(Address::Heap(11)),
            Cell::Functor(String::from("b"), 0),
            Cell::Structure(Address::Heap(13)),
            Cell::Functor(String::from("f"), 2),
            Cell::Structure(Address::Heap(11)),
            Cell::Variable(Address::Heap(15)),
        ];
        
        assert_eq!(inst.heap, expected_heap);

        inst.unify(Address::Heap(6), Address::Heap(12));
        
        let expected_heap = vec![
            Cell::Structure(Address::Heap(1)),
            Cell::Functor(String::from("a"), 0),
            Cell::Structure(Address::Heap(3)),
            Cell::Functor(String::from("g"), 2),
            Cell::Variable(Address::Heap(14)),      // 4 -> 14
            Cell::Structure(Address::Heap(1)),
            Cell::Structure(Address::Heap(7)),
            Cell::Functor(String::from("f"), 2),
            Cell::Variable(Address::Heap(4)),
            Cell::Structure(Address::Heap(3)),
            Cell::Structure(Address::Heap(11)),
            Cell::Functor(String::from("b"), 0),
            Cell::Structure(Address::Heap(13)),
            Cell::Functor(String::from("f"), 2),
            Cell::Structure(Address::Heap(11)),
            Cell::Variable(Address::Heap(9)),       // 15 -> 9
        ];

        assert_eq!(inst.heap, expected_heap);

    }

    #[test]
    fn program_instructions() {
        let mut inst = Instance::new();
        inst.load(instruction_set_1());
        inst.load(instruction_set_2());
        inst.execute();

        let expected_heap = vec![
            Cell::Structure(Address::Heap(1)),
            Cell::Functor(String::from("h"), 2),
            Cell::Variable(Address::Heap(12)),
            Cell::Variable(Address::Heap(14)),
            Cell::Structure(Address::Heap(5)),
            Cell::Functor(String::from("f"), 1),
            Cell::Variable(Address::Heap(3)),
            Cell::Structure(Address::Heap(8)),
            Cell::Functor(String::from("p"), 3),
            Cell::Variable(Address::Heap(2)),
            Cell::Structure(Address::Heap(1)),
            Cell::Structure(Address::Heap(5)),
            Cell::Structure(Address::Heap(13)),
            Cell::Functor(String::from("f"), 1),
            Cell::Variable(Address::Heap(15)),
            Cell::Structure(Address::Heap(16)),
            Cell::Functor(String::from("a"), 0),
        ];

        assert_eq!(inst.heap, expected_heap);

    }
    
    #[test]
    fn argument_instructions_1() {
        let mut inst_1 = Instance::new();
        inst_1.load(instruction_set_1());
        inst_1.execute();

        let mut inst_2 = Instance::new();
        inst_2.load(instruction_set_3());
        let (label, instructions) = instruction_set_4();
        inst_2.load_labeled(label, instructions);
        inst_2.execute();

        assert_eq!(inst_1.heap, inst_2.heap);
    }
}