From dc2d96c4fff082740cf45a1c9731ea76811f451b Mon Sep 17 00:00:00 2001
From: Ben Ruijl <benruyl@gmail.com>
Date: Wed, 6 Nov 2024 14:24:20 +0100
Subject: [PATCH] Handle self-loops in spanning tree representation

- Limit atom sizes in cache and drop them upon exiting thread
---
 src/atom/representation.rs | 13 +++++++++++++
 src/graph.rs               |  7 ++++++-
 src/state.rs               | 14 +++++---------
 3 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/src/atom/representation.rs b/src/atom/representation.rs
index 5fb92de..82a393f 100644
--- a/src/atom/representation.rs
+++ b/src/atom/representation.rs
@@ -212,6 +212,19 @@ impl Atom {
             _ => unreachable!("Unknown type {}", raw[0]),
         }
     }
+
+    /// Get the capacity of the underlying buffer.
+    pub(crate) fn get_capacity(&self) -> usize {
+        match self {
+            Atom::Num(n) => n.data.capacity(),
+            Atom::Var(v) => v.data.capacity(),
+            Atom::Fun(f) => f.data.capacity(),
+            Atom::Mul(m) => m.data.capacity(),
+            Atom::Add(a) => a.data.capacity(),
+            Atom::Pow(p) => p.data.capacity(),
+            Atom::Zero => 0,
+        }
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
diff --git a/src/graph.rs b/src/graph.rs
index a3af523..36ef2fa 100644
--- a/src/graph.rs
+++ b/src/graph.rs
@@ -193,7 +193,8 @@ impl SpanningTree {
                 let node = self.nodes[n].back_edges[back_edge_index];
                 back_edge_index += 1;
 
-                if self.nodes[node].parent == n {
+                if node == n {
+                    // self-loop
                     continue;
                 }
 
@@ -390,6 +391,10 @@ impl<N, E> Graph<N, E> {
                     edge.vertices.0
                 };
 
+                if n == target {
+                    tree_nodes[n].back_edges.push(n);
+                }
+
                 if tree_nodes[target].position.is_none() {
                     nodes_to_visit.push((target, n));
                 }
diff --git a/src/state.rs b/src/state.rs
index 759f396..547483b 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -1,7 +1,6 @@
 use byteorder::{ReadBytesExt, WriteBytesExt};
 use std::hash::Hash;
 use std::io::{Read, Write};
-use std::mem::ManuallyDrop;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, RwLock};
 use std::thread::LocalKey;
@@ -73,11 +72,8 @@ static VARIABLE_LISTS: AppendOnlyVec<Arc<Vec<Variable>>> = AppendOnlyVec::new();
 static SYMBOL_OFFSET: AtomicUsize = AtomicUsize::new(0);
 
 thread_local!(
-    /// A thread-local workspace, that stores recyclable atoms. By making it const and
-    /// `ManuallyDrop`, the fastest implementation is chosen for the current platform.
-    /// In principle this leaks memory, but Symbolica only uses thread pools that live as
-    /// long as the main thread, so this is no issue.
-    static WORKSPACE: ManuallyDrop<Workspace> = const { ManuallyDrop::new(Workspace::new()) }
+    /// A thread-local workspace, that stores recyclable atoms.
+    static WORKSPACE: Workspace = const { Workspace::new() }
 );
 
 /// A global state, that stores mappings from variable and function names to ids.
@@ -713,7 +709,7 @@ pub struct Workspace {
 
 impl Workspace {
     const ATOM_BUFFER_MAX: usize = 30;
-    const ATOM_CACHE_SIZE_MAX: usize = 10_000_000;
+    const ATOM_CACHE_SIZE_MAX: usize = 20_000_000;
 
     /// Create a new workspace.
     const fn new() -> Self {
@@ -724,7 +720,7 @@ impl Workspace {
 
     /// Get a thread-local workspace.
     #[inline]
-    pub fn get_local() -> &'static LocalKey<ManuallyDrop<Workspace>> {
+    pub fn get_local() -> &'static LocalKey<Workspace> {
         LicenseManager::check();
 
         &WORKSPACE
@@ -848,7 +844,7 @@ impl Drop for RecycledAtom {
             return;
         }
 
-        if self.0.as_view().get_byte_size() > Workspace::ATOM_CACHE_SIZE_MAX {
+        if self.0.get_capacity() > Workspace::ATOM_CACHE_SIZE_MAX {
             return;
         }