Need help implementing drop for tree

I have this Tree structure, which holds parent via Weak and childs via vector of Rc:

pub type NodePtr = Rc<RefCell<Tree>>;
pub type WeakNodePtr = Weak<RefCell<Tree>>;

pub struct Tree {
    pub data: HashMap<String, NodePtr>,
    pub parent: Option<WeakNodePtr>,
    pub childs: Vec<NodePtr>,
}

Currently I can generate very large trees, but can't drop those, because recursion blows the stack. I know how traverse this tree in iterative way:

pub fn iterative_print(node: &NodePtr) {
    if !node.borrow().childs.is_empty() {
        let mut stack = vec![];
        stack.push(Tree::get_parent(node.borrow().childs.last().expect("no childs")));
        while !stack.is_empty() {
            let current = stack.pop().expect("empty").expect("no parent");
            println!("{:?}", current.borrow().data);
            for child in current.borrow().childs.iter() {
                if !child.borrow().childs.is_empty() {
                    stack.push(Tree::get_parent(child.borrow().childs.last().expect("no childs")));
                } else {
                    println!("{:?}", child.borrow().data)
                }
            }
        }
    }
}

where Tree::get_parent is:

pub fn get_parent(node: &NodePtr) -> Option<NodePtr> {
    Weak::upgrade(node.borrow().parent.as_ref()?)
}

I suppose that I should convert that iterative_print procedure to the drop where instead of printing the node that has no childs I free the memory and pop the stack. This is what I've got:

impl Drop for Tree {
    fn drop(&mut self) {
        for c in self.childs.iter() {
            let mut stack = vec![];
            stack.push(c.clone());
            while !stack.is_empty() {
                let current = stack.pop().expect("empty");
                for child in current.borrow().childs.iter() {
                    if !child.borrow().childs.is_empty() {
                        stack.push(Tree::get_parent(
                            child.borrow().childs.last().expect("no childs 2"),
                        ).expect("no_parent"));
                        continue
                    } else {
                        std::mem::drop(child);
                    }
                }
                std::mem::drop(current);
            }
        }
    }
}

Checking with valgrind I see no leaks, but Dropping is very slow, like if I'm doing extra work so I'm suspicious that it doesn't do things correctly. Especially when it comes to reference counting. I also suspect that because I'm cloning references to the stack I actually increase the life time of Rc here, so this code is all wrong. Any suggestions?

One questions: Why are you calling get_parent on an element of child.childs, wouldn't that always return child? (for a valid tree)

Would this work?

impl Drop for Tree {
    fn drop(&mut self) {
        let mut stack = vec![]; // reuse the stack allocation
        
        for c in self.childs.iter().cloned() {
            stack.push(c); // initialize the stack
            
            // keep popping off of the stack till it is empty
            while let Some(current) = stack.pop() {
                let mut current = current.borrow_mut();

                // extend the stack with the children, so we can handle them later
                stack.extend_from_slice(&current.childs);

                // prevent current from recursing on drop, 
                // because we will clean up the children here
                current.childs.clear();

                // current will be automatically dropped here
            }
        }
    }
}

oh, perhaps my mind exhausted already for today and I've wrote complete gibberish. Yeah, this seem to work for the tree with depth of 1 million nodes.

1 Like

this uses unstable feature, and I'm writing in stable branch.

Also it seems that your previous drop breaks my code. Perhaps it clears too aggressively? The one that is generated by the compiler doesn't break my code.

heh, sorry about that I normally use nightly so I didn't notice

Your right, fixed that here

impl Drop for Tree {
    fn drop(&mut self) {
        // reuse the stack allocation, take all children
        let mut stack = std::mem::replace(&mut self.childs, Vec::new());
        
        // keep popping off of the stack till it is empty
        while let Some(mut current) = stack.pop() {
            // only take the children if `current` `Tree` is going to be dropped
            // right now if there are outstanding `Rc`s elsewhere,
            // don't clean this up just yet
            if let Some(current) = Rc::get_mut(&mut current) {
                let current = RefCell::get_mut(current);
                
                // extend the stack with the children, so we can handle them later
                stack.extend_from_slice(&current.childs);
    
                // prevent current from recursing on drop, 
                // because we will clean up the children here
                current.childs.clear();
            }
            
            // current will be automatically dropped here
        }
    }
}

std::mem::take has stablized in 1.40.0, are you using an older version?

In any case std::mem::take(&mut value) is exactly the same as std::mem::replace(&mut value, Default::default()), so easy replacement

it seems it's recursing.I'm gedding stack overflow with this drop, on big trees. Will attach playground link, with minimal Tree in a sec.

oh you're right, yes, I've forgot to update :slight_smile:

I think you are storing Rc's elsewhere and those are the problem, not the drop implementation itself. The latest drop impl that I gave will drop every child that is currently available to be dropped (i.e. isn't being shared somewhere else).

IIUUC you mean I had memory leaked with drop generated by compiler?

No, I mean something like this

let tree = Tree::new();
// initialize tree

let child: NodePtr = get_some_child(&mut tree);

drop(tree); // this won't drop `child`

drop(child);

But thinking about it a bit more, I don't see how this would turn into a stack overflow

I'll take a look at the playground example, once you post it

playground

Actually I've figured that removing the clone() in line 50 fixes stack overflow. Can you explain why? Clippy says that this clone is #[warn(clippy::redundant_clone)], but I wouldn't imagine that it will cause stack overflow.

Ahhhhhh, I forgot that Rc::get_mut will return None if there are any weak pointers active, but we don't care about weak pointers (because we are about to drop the last Rc which will invalidate all weak pointers anyways)!

Since children always have a weak pointer to their parent, this fools Rc::get_mut and turns the old drop into the recursive drop. But we can just check this Rc::strong_count directly to see if we are going to drop, and handle that accordingly.

impl Drop for Tree {
    fn drop(&mut self) {
        let mut stack = std::mem::replace(&mut self.childs, Vec::new());
        while let Some(current) = stack.pop() {
            if Rc::strong_count(&current) == 1 {
                let mut current = current.borrow_mut();
                stack.extend_from_slice(&current.childs);
                current.childs.clear();
            }
        }
    }
}

The reason why removing clone worked is because you were then dropping the parents as soon as possible, every time you overwrote tmp, not at the end of main. Adding the clone forced the full recursive drop at the end of main.

1 Like

Yes this works fine, thanks a lot for explanations!

1 Like