From 3a7de80db98c53fcadd73660d9a3f189b62e4d5b Mon Sep 17 00:00:00 2001
From: Josef Podany <jsfpdn@gmail.com>
Date: Sun, 13 Oct 2024 18:27:26 +0200
Subject: [PATCH] vtree: retrieve vtrees for literals

---
 sdd-rs-lib/vtree.rs | 58 ++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 57 insertions(+), 1 deletion(-)

diff --git a/sdd-rs-lib/vtree.rs b/sdd-rs-lib/vtree.rs
index 6d90ffa..e8ac1bc 100644
--- a/sdd-rs-lib/vtree.rs
+++ b/sdd-rs-lib/vtree.rs
@@ -2,7 +2,7 @@ use std::{cell::RefCell, fmt::Debug, rc::Rc};
 
 use crate::{
     dot_writer::{Dot, DotWriter, Edge, NodeType},
-    literal::VarLabel,
+    literal::{Literal, VarLabel},
     manager::SddManager,
 };
 
@@ -366,6 +366,25 @@ impl VTreeManager {
         VTreeManager::set_inorder_indices(self.root.clone().unwrap(), 0);
     }
 
+    pub(crate) fn get_variable_vtree(&self, variable: &VarLabel) -> Option<VTreeRef> {
+        fn find_vtree(vtree: &VTreeRef, variable: &VarLabel) -> Option<VTreeRef> {
+            match vtree.borrow().node.clone() {
+                Node::Internal(lc, rc) => find_vtree(&lc, variable)
+                    .or(find_vtree(&rc, variable))
+                    .or(None),
+                Node::Leaf(candidate_variable) => {
+                    if *variable == candidate_variable {
+                        Some(vtree.clone())
+                    } else {
+                        None
+                    }
+                }
+            }
+        }
+
+        find_vtree(&self.root.clone().unwrap(), variable)
+    }
+
     fn get_vtree(&self, index: u16) -> Option<VTreeRef> {
         // TODO: This will get obsolete once VTrees are stored in a single hashmap.
         let Some(mut current) = self.root.clone() else {
@@ -782,4 +801,41 @@ mod test {
         assert_eq!(ord, VTreeOrder::Inequal);
         assert_eq!(lca.borrow().idx, root_idx);
     }
+
+    #[test]
+    fn literal_indices() {
+        let var_label_index = |vtree: Option<VTreeRef>| -> u16 { vtree.unwrap().borrow().idx };
+
+        let mut manager = VTreeManager::new();
+        manager.add_variable(VarLabel::new("A"));
+        manager.add_variable(VarLabel::new("B"));
+        manager.add_variable(VarLabel::new("C"));
+        manager.add_variable(VarLabel::new("D"));
+        //     1
+        //   /   \
+        //  0     3
+        //  A   /   \
+        //     2     5
+        //     B   /   \
+        //        4     6
+        //        C     D
+
+        assert_eq!(
+            var_label_index(manager.get_variable_vtree(&VarLabel::new("A"))),
+            0
+        );
+        assert_eq!(
+            var_label_index(manager.get_variable_vtree(&VarLabel::new("B"))),
+            2
+        );
+        assert_eq!(
+            var_label_index(manager.get_variable_vtree(&VarLabel::new("C"))),
+            4
+        );
+        assert_eq!(
+            var_label_index(manager.get_variable_vtree(&VarLabel::new("D"))),
+            6
+        );
+        assert_eq!(manager.get_variable_vtree(&VarLabel::new("E")), None);
+    }
 }