Skip to main content

prolog2/heap/
query_heap.rs

1use crate::heap::{heap::Tag, symbol_db::SymbolDB};
2use fsize::fsize;
3use std::{
4    mem,
5    ops::{Index, IndexMut, Range},
6    sync::atomic::{AtomicUsize, Ordering::Acquire},
7};
8
9use super::heap::{Cell, Heap};
10
11static HEAP_ID_COUNTER: AtomicUsize = AtomicUsize::new(1);
12
13/// Working heap for proof search.
14///
15/// Wraps a shared read-only program heap (`&[Cell]`) and an
16/// owned mutable cell buffer for query-time allocations. Supports
17/// branching via an optional parent pointer for backtracking.
18pub struct QueryHeap<'a> {
19    id: usize,
20    pub(crate) cells: Vec<Cell>,
21    prog_cells: &'a [Cell],
22    // TODO: handle branching query heap multi-threading
23    root: Option<*const QueryHeap<'a>>,
24}
25
26impl<'a> QueryHeap<'a> {
27    pub fn new(prog_cells: &'a [Cell], root: Option<*const QueryHeap<'a>>) -> QueryHeap<'a> {
28        let id = HEAP_ID_COUNTER.fetch_add(1, Acquire);
29        QueryHeap {
30            id,
31            cells: Vec::new(),
32            prog_cells,
33            root,
34        }
35    }
36
37    pub fn branch(&self, count: usize) -> Vec<QueryHeap<'a>> {
38        let mut branch_heap = Vec::with_capacity(count);
39        for _ in 0..count {
40            branch_heap.push(QueryHeap::new(self.prog_cells, Some(self)));
41        }
42        branch_heap
43    }
44
45    fn get_symbol_db_id(&self, addr: usize) -> usize {
46        if addr < self.prog_cells.len() {
47            0
48        } else {
49            self.id
50        }
51    }
52}
53
54impl Heap for QueryHeap<'_> {
55    #[inline(always)]
56    fn deref_addr(&self, mut addr: usize) -> usize {
57        loop {
58            let cell = self[addr];
59            match cell {
60                (Tag::Ref, pointer) if addr == pointer => return pointer,
61                (Tag::Ref, pointer) => addr = pointer,
62                _ => return addr,
63            }
64        }
65    }
66
67    fn heap_push(&mut self, cell: Cell) -> usize {
68        let i = self.heap_len();
69        self.cells.push(cell);
70        i
71    }
72
73    fn heap_len(&self) -> usize {
74        match self.root {
75            Some(root) => unsafe { &*root }.heap_len() + self.cells.len(),
76            None => self.prog_cells.len() + self.cells.len(),
77        }
78    }
79
80    fn get_id(&self) -> usize {
81        self.id
82    }
83
84    /** Create String to represent cell, can be recursively used to format complex structures or list */
85    fn term_string(&self, addr: usize) -> String {
86        // println!("[{addr}]:{:?}", self[addr]);
87        let addr = self.deref_addr(addr);
88        match self[addr].0 {
89            Tag::Con => SymbolDB::get_const(self[addr].1).to_string(),
90            Tag::Func => self.func_string(addr),
91            Tag::Lis => self.list_string(addr),
92            Tag::ELis => "[]".into(),
93            Tag::Arg => match SymbolDB::get_var(addr, self.get_symbol_db_id(addr)) {
94                Some(symbol) => symbol.to_string(),
95                None => format!("Arg_{}", self[addr].1),
96            },
97            Tag::Ref => match SymbolDB::get_var(self.deref_addr(addr), self.get_symbol_db_id(addr))
98                .to_owned()
99            {
100                Some(symbol) => symbol.to_string(),
101                None => format!("Ref_{}", self[addr].1),
102            },
103            Tag::Int => {
104                let value: isize = unsafe { mem::transmute_copy(&self[addr].1) };
105                format!("{value}")
106            }
107            Tag::Flt => {
108                let value: fsize = unsafe { mem::transmute_copy(&self[addr].1) };
109                format!("{value}")
110            }
111            Tag::Tup => self.tuple_string(addr),
112            Tag::Set => self.set_string(addr),
113            Tag::Str => self.term_string(self[addr].1),
114            Tag::Stri => SymbolDB::get_string(self[addr].1).to_string(),
115        }
116    }
117
118    fn truncate(&mut self, mut len: usize) {
119        if len < self.prog_cells.len() {
120            panic!("Can't truncate to prog cells")
121        }
122        len -= self.prog_cells.len();
123        self.cells.resize(len, (Tag::Ref, 0));
124    }
125}
126
127impl Index<usize> for QueryHeap<'_> {
128    type Output = Cell;
129
130    fn index(&self, index: usize) -> &Self::Output {
131        if index < self.prog_cells.len() {
132            &self.prog_cells[index]
133        } else {
134            if let Some(root) = self.root {
135                let root = unsafe { &*root };
136                if index < root.heap_len() {
137                    // Index is in root's query cells
138                    &root[index]
139                } else {
140                    // Index is in our own cells
141                    &self.cells[index - root.heap_len()]
142                }
143            } else {
144                &self.cells[index - self.prog_cells.len()]
145            }
146        }
147    }
148}
149
150impl IndexMut<usize> for QueryHeap<'_> {
151    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
152        if index < self.prog_cells.len() {
153            panic!("Can't get mutable reference to program heap cell");
154        } else {
155            if let Some(root) = self.root {
156                let root = unsafe { &*root };
157                let root_heap_len = root.heap_len(); // prog_cells.len() + root.cells.len()
158                if index < root_heap_len {
159                    // Index is in root's query cells - deny mutable access
160                    panic!("Can't get mutable reference to root heap cell");
161                } else {
162                    // Index is in our own cells
163                    &mut self.cells[index - root_heap_len]
164                }
165            } else {
166                &mut self.cells[index - self.prog_cells.len()]
167            }
168        }
169    }
170}
171
172impl Index<Range<usize>> for QueryHeap<'_> {
173    type Output = [Cell];
174
175    fn index(&self, index: Range<usize>) -> &Self::Output {
176        let len = self.prog_cells.len();
177
178        if index.start < len && index.end < len {
179            &self.prog_cells[index]
180        } else if index.start >= len && self.root.is_none() {
181            &self.cells[index.start - len..index.end - len]
182        } else {
183            panic!("Range splits static and mutable heap")
184        }
185    }
186}