changes

src/main.rs

@@ -59,65 +59,74 @@ struct RoyaltyTree {
     tree: Vec<Royalty>,
     }
 
+struct Royalty {
+    weight: FiNum,      // Here and below
+    lazy: FiNum,	// To be distributed to here and to below based on weights
+    acc: FiNum,         // Accumulated here
+    }
+
+impl Royalty {
+    fn new() -> Self {
+        Royalty { weight:FiNum::zero(), lazy:FiNum::zero(), acc:FiNum::zero() }
+        }
+    }
+
 impl RoyaltyTree {
     fn new() -> Self {
         RoyaltyTree { tree:Vec::new() }
 	}
-    fn get_weight(&self, index:usize) -> FiNum {
-        println!("Get_weight {} self.tree.len() {} ",index,self.tree.len());
-        if index<self.tree.len() { self.tree[index].weight }
-	else {
+    fn weight_here_below(&self, index:usize) -> FiNum {
+        if index<self.tree.len() { self.tree[index].weight } else { FiNum::zero() }
+	}
+    fn weight_below(&self, index:usize) -> FiNum {
+        if index&1==0 {
+	    FiNum::zero()
+	} else {
+	    self.weight_here_below(wt_left (index).unwrap())+self.weight_here_below(wt_right(index).unwrap())
+	    }
+	}
+    fn weight_here(&self, index:usize) -> FiNum {
+        if index&1==0 {
+	    self.weight_here_below(index)
+	} else {
+	    self.weight_here_below(index)-self.weight_below(index)
+	    }
+	}
+    fn expand_to(&mut self, index: usize) -> &mut Self {
+        for _ in self.tree.len()..=index+1 { self.tree.push(Royalty::new()) }
+	self
+        }
+    fn capture0(&mut self, index: usize) -> &mut Self {
+        if index<self.tree.len() {
+	    if index&1==0 {
+	        self.tree[index].acc+=self.tree[index].lazy;
+		self.tree[index].lazy=FiNum::zero();
+	    } else {
+	        let lazy=self.tree[index].lazy;
+    	        let d1=lazy*self.weight_here(index)/self.weight_here_below(index);
+		let d02=lazy-d1;
+		let index_left =wt_left (index).unwrap();
+		let index_right=wt_right(index).unwrap();
+		let d0=d02*self.weight_here_below(index_left)/self.weight_below(index);
+		let d2=d02-d0;
+		assert!(lazy==d0+d1+d2,"Distributing amounts that don't add up.");
+		self.expand_to(index_right);
+		self.tree[index_left ].lazy+=d0;
+		self.tree[index_right].lazy+=d2;
+		self.tree[index].acc+=d1;
+		self.tree[index].lazy=FiNum::zero();
+		}
+	    }
+	self
+        }
+    fn capture(&mut self, index: usize) -> &mut Self { // Clear out lazy by drilling from root to index
         let ff=wt_forefather(index);
 	let mut index=index;
         while index!=ff {
-	        if index<self.tree.len() { return self.tree[index].weight }
-		else { index=wt_parent(index) }
+	    self.capture0(index);
+	    index=wt_parent(index);
 	    }
-	    while index>=self.tree.len() { index=wt_left(index).unwrap() }
-	    self.tree[index].weight
-	    }
-	}
-    }
-
-struct Royalty {
-    weight: FiNum,      // Here and below
-    acc: FiNum,         // Here and Below
-    }
-
-impl Royalty {
-    fn new(weight: FiNum) -> Self {
-        Royalty { weight:weight, acc:FiNum::new(0u64) }
-        }
-    }
-
-impl RoyaltyTree {
-    fn insert(&mut self, weight: FiNum, royalty:FiNum) -> usize {
-        let last=self.tree.len();
-	let mut nweight=weight;
-	if let Some(left )=wt_left (last) { if left <self.tree.len() { nweight+=self.tree[left ].weight } }
-	if let Some(right)=wt_right(last) { if right<self.tree.len() { nweight+=self.tree[right].weight } }
-        self.tree.push(Royalty::new(nweight));
-        let forefather=wt_forefather(last);
-        let mut pivot=last;
-        while pivot!=forefather {
-            pivot=wt_parent(pivot);
-	    if pivot<self.tree.len() { self.tree[pivot].weight+=weight }
-            }
-	self.tree[forefather].acc+=royalty;
-	last
-        }
-    fn capture(&mut self, index:usize, weight: FiNum) -> FiNum {
-        let ff=wt_forefather(index);
-	let mut index=index;
-	let mut rval=FiNum::zero();
-	loop {
-	    let cap=self.tree[index].acc*weight/self.tree[index].weight;
-	    self.tree[index].weight-=weight;
-	    self.tree[index].acc-=cap;
-	    rval+=cap;
-	    if index!=ff { index=wt_parent(index); } else { break; }
-	    }
-	rval
+	self
         }
     fn dump(&self) {
         for index in 0..self.tree.len() {