From 1653e75b6bbc3c440f27ca34b2929a64edfc2a55 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Tue, 25 Feb 2020 16:05:59 -0800
Subject: [PATCH] feat: maxsharing for arrays and strings

---
 src/runtime/maxsharing.cpp     | 30 +++++++++++++++++++++-----
 tests/lean/run/maxsharing.lean | 39 +++++++++++++++++++++++++++++++---
 2 files changed, 61 insertions(+), 8 deletions(-)

diff --git a/src/runtime/maxsharing.cpp b/src/runtime/maxsharing.cpp
index 546f8b3d7e..045b1da3d0 100644
--- a/src/runtime/maxsharing.cpp
+++ b/src/runtime/maxsharing.cpp
@@ -223,9 +223,22 @@ class max_sharing_fn {
     }
 
     void visit_array(b_obj_arg a) {
-        // TODO(Leo)
-        lean_inc(a);
-        save(a, a);
+        clear_children();
+        bool missing_children = false;
+        size_t sz = array_size(a);
+        for (size_t i = 0; i < sz; i++) {
+            if (!push_child(lean_array_get_core(a, i))) {
+                missing_children = true;
+            }
+        }
+        if (missing_children)
+            return;
+        lean_array_object * new_a = (lean_array_object*)lean_alloc_array(sz, sz);
+        for (size_t i = 0; i < sz; i++) {
+            lean_inc(m_children[i]);
+            lean_array_set_core((lean_object*)new_a, i, m_children[i]);
+        }
+        save(a, (lean_object*)new_a);
     }
 
     void visit_sarray(b_obj_arg a) {
@@ -235,8 +248,15 @@ class max_sharing_fn {
     }
 
     void visit_string(b_obj_arg a) {
-        lean_inc(a);
-        save(a, a);
+        size_t sz     = lean_string_size(a);
+        size_t len    = lean_string_len(a);
+        lean_string_object * new_a = (lean_string_object*)lean_alloc_string(sz, sz, len);
+        lean_set_st_header((lean_object*)new_a, LeanString, 0);
+        new_a->m_size     = sz;
+        new_a->m_capacity = sz;
+        new_a->m_length   = len;
+        memcpy(new_a->m_data, lean_to_string(a)->m_data, sz);
+        save(a, (lean_object*)new_a);
     }
 
     void visit_ctor(b_obj_arg a) {
diff --git a/tests/lean/run/maxsharing.lean b/tests/lean/run/maxsharing.lean
index ca888e59ab..823e3300b8 100644
--- a/tests/lean/run/maxsharing.lean
+++ b/tests/lean/run/maxsharing.lean
@@ -73,14 +73,47 @@ check $ ptrAddrUnsafe x != ptrAddrUnsafe y;
 let s := MaxSharing.State.empty;
 let (x, s) := s.maxSharing x;
 let (y, s) := s.maxSharing y;
--- check $ ptrAddrUnsafe x == ptrAddrUnsafe y;
+check $ ptrAddrUnsafe x == ptrAddrUnsafe y;
 let (z, s) := s.maxSharing ["world"];
 let (x, s) := s.maxSharing x;
--- check $ ptrAddrUnsafe x == ptrAddrUnsafe y;
+check $ ptrAddrUnsafe x == ptrAddrUnsafe y;
 check $ ptrAddrUnsafe x != ptrAddrUnsafe z;
 IO.println x;
 IO.println y;
 IO.println z
 
+#eval tst4
 
-#eval tst3
+@[noinline] def mkList1 (x : Nat) : List Nat := List.replicate x x
+@[noinline] def mkList2 (x : Nat) : List Nat := List.replicate x x
+@[noinline] def mkArray1 (x : Nat) : Array (List Nat) :=
+#[ mkList1 x, mkList2 x, mkList2 (x+1) ]
+@[noinline] def mkArray2 (x : Nat) : Array (List Nat) :=
+mkArray1 x
+
+unsafe def tst5 : IO Unit := do
+let a := mkArray1 3;
+let b := mkArray2 3;
+let c := mkArray2 4;
+IO.println a;
+IO.println b;
+IO.println c;
+check $ ptrAddrUnsafe a != ptrAddrUnsafe b;
+check $ ptrAddrUnsafe a != ptrAddrUnsafe c;
+check $ ptrAddrUnsafe (a.get! 0) != ptrAddrUnsafe (a.get! 1);
+check $ ptrAddrUnsafe (a.get! 0) != ptrAddrUnsafe (a.get! 2);
+check $ ptrAddrUnsafe (b.get! 0) != ptrAddrUnsafe (b.get! 1);
+check $ ptrAddrUnsafe (c.get! 0) != ptrAddrUnsafe (c.get! 1);
+let s := MaxSharing.State.empty;
+let (a, s) := s.maxSharing a;
+let (b, s) := s.maxSharing b;
+let (c, s) := s.maxSharing c;
+check $ ptrAddrUnsafe a == ptrAddrUnsafe b;
+check $ ptrAddrUnsafe a != ptrAddrUnsafe c;
+check $ ptrAddrUnsafe (a.get! 0) == ptrAddrUnsafe (a.get! 1);
+check $ ptrAddrUnsafe (a.get! 0) != ptrAddrUnsafe (a.get! 2);
+check $ ptrAddrUnsafe (b.get! 0) == ptrAddrUnsafe (b.get! 1);
+check $ ptrAddrUnsafe (c.get! 0) == ptrAddrUnsafe (c.get! 1);
+pure ()
+
+#eval tst5