a.diff

diff --git a/Python/tracemalloc.c b/Python/tracemalloc.c
index 66de957..65273c1 100644
--- a/Python/tracemalloc.c
+++ b/Python/tracemalloc.c
@@ -12,7 +12,7 @@
 #include "pycore_traceback.h"     // _Py_DumpASCII()
 #include "pycore_tstate.h"        // _PyThreadStateImpl
 
-#include <math.h>                 // log()
+#include <math.h>                 // log(), expm1()
 #include <stdlib.h>               // malloc()
 
 #define tracemalloc_config _PyRuntime.tracemalloc.config
@@ -185,6 +185,39 @@ new_sampling_seed(void)
 }
 
 
+/* Horvitz-Thompson weight for a sampled allocation: the allocation's true
+   size divided by its sampling probability p = 1 - exp(-size/interval).
+
+   This makes the per-trace size an unbiased estimator of the bytes the
+   sample represents, for any allocation size relative to the interval.
+   Crediting the bytes accumulated since the last sample instead would
+   over-attribute allocations that are comparable to or larger than the
+   interval (they sweep up the small allocations pending in the counter)
+   and correspondingly under-attribute the small-allocation tail.
+
+   expm1() keeps p accurate when size << interval, where 1 - exp(-x) would
+   lose nearly all its significant digits to cancellation.  The weight is
+   always at least the real size and is clamped to SIZE_MAX. */
+static size_t
+sample_weight(size_t byte_size, size_t interval)
+{
+    if (byte_size == 0) {
+        return 0;
+    }
+    double p = -expm1(-(double)byte_size / (double)interval);
+    if (p <= 0.0) {
+        /* Defensive: p is mathematically in (0, 1) for byte_size > 0. */
+        return byte_size;
+    }
+    double weight = (double)byte_size / p;
+    if (weight >= (double)SIZE_MAX) {
+        return SIZE_MAX;
+    }
+    size_t w = (size_t)weight;
+    return w >= byte_size ? w : byte_size;
+}
+
+
 /* Seed the PRNG and draw the first threshold.  Called on a thread's first
    sampled allocation (prng_state == 0 after zero-init or explicit reset). */
 static void
@@ -259,7 +292,7 @@ new_global_threshold(size_t interval)
 /* Sampling decision.  Returns true if this allocation should be traced,
    false if it should be skipped.  When returning true,
    *effective_size is set: to byte_size in exact mode (interval == 0),
-   or to the upscaled Poisson weight in sampled mode. */
+   or to the Horvitz-Thompson weight (see sample_weight) in sampled mode. */
 static bool
 should_sample(size_t byte_size, size_t *effective_size)
 {
@@ -279,7 +312,7 @@ should_sample(size_t byte_size, size_t *effective_size)
         if (s->bytes_since_last_sample < s->threshold) {
             return false;
         }
-        *effective_size = s->bytes_since_last_sample;
+        *effective_size = sample_weight(byte_size, interval);
         s->bytes_since_last_sample = 0;
         s->threshold = new_sample_threshold(interval, &s->prng_state);
         return true;
@@ -337,7 +370,7 @@ should_sample(size_t byte_size, size_t *effective_size)
     if (!sampled) {
         return false;
     }
-    *effective_size = (total > (uint64_t)SIZE_MAX) ? SIZE_MAX : (size_t)total;
+    *effective_size = sample_weight(byte_size, interval);
     return true;
 }