From 0971e1e4e33edf1cd0d5be808d1eb092ffeab9f3 Mon Sep 17 00:00:00 2001
From: David Goldblatt <davidgoldblatt@fb.com>
Date: Mon, 30 Nov 2020 14:34:27 -0800
Subject: [PATCH] hpdata: Use addr/size instead of begin/npages.

This is easier for the users of the hpdata.
---
 include/jemalloc/internal/hpdata.h |  4 ++--
 src/hpdata.c                       | 16 ++++++++++++----
 src/psset.c                        | 20 ++++++--------------
 3 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/include/jemalloc/internal/hpdata.h b/include/jemalloc/internal/hpdata.h
index d221c577..cb034eae 100644
--- a/include/jemalloc/internal/hpdata.h
+++ b/include/jemalloc/internal/hpdata.h
@@ -133,7 +133,7 @@ void hpdata_init(hpdata_t *hpdata, void *addr, uint64_t age);
  * Given an hpdata which can serve an allocation request, pick and reserve an
  * offset within that allocation.
  */
-size_t hpdata_reserve_alloc(hpdata_t *hpdata, size_t npages);
-void hpdata_unreserve(hpdata_t *hpdata, size_t start, size_t npages);
+void *hpdata_reserve_alloc(hpdata_t *hpdata, size_t sz);
+void hpdata_unreserve(hpdata_t *hpdata, void *begin, size_t sz);
 
 #endif /* JEMALLOC_INTERNAL_HPDATA_H */
diff --git a/src/hpdata.c b/src/hpdata.c
index a876a302..847eb9da 100644
--- a/src/hpdata.c
+++ b/src/hpdata.c
@@ -28,8 +28,10 @@ hpdata_init(hpdata_t *hpdata, void *addr, uint64_t age) {
 	fb_init(hpdata->active_pages, HUGEPAGE_PAGES);
 }
 
-size_t
-hpdata_reserve_alloc(hpdata_t *hpdata, size_t npages) {
+void *
+hpdata_reserve_alloc(hpdata_t *hpdata, size_t sz) {
+	assert((sz & PAGE_MASK) == 0);
+	size_t npages = sz >> LG_PAGE;
 	assert(npages <= hpdata_longest_free_range_get(hpdata));
 
 	size_t result;
@@ -91,11 +93,17 @@ hpdata_reserve_alloc(hpdata_t *hpdata, size_t npages) {
 	}
 	hpdata_longest_free_range_set(hpdata, largest_unchosen_range);
 
-	return result;
+	return (void *)(
+	    (uintptr_t)hpdata_addr_get(hpdata) + (result << LG_PAGE));
 }
 
 void
-hpdata_unreserve(hpdata_t *hpdata, size_t begin, size_t npages) {
+hpdata_unreserve(hpdata_t *hpdata, void *addr, size_t sz) {
+	assert((sz & PAGE_MASK) == 0);
+	size_t begin = ((uintptr_t)addr - (uintptr_t)hpdata_addr_get(hpdata))
+	    >> LG_PAGE;
+	assert(begin < HUGEPAGE_PAGES);
+	size_t npages = sz >> LG_PAGE;
 	size_t old_longest_range = hpdata_longest_free_range_get(hpdata);
 
 	fb_unset_range(hpdata->active_pages, HUGEPAGE_PAGES, begin, npages);
diff --git a/src/psset.c b/src/psset.c
index 2e3558c1..c31520fd 100644
--- a/src/psset.c
+++ b/src/psset.c
@@ -152,10 +152,8 @@ psset_fit(psset_t *psset, size_t size) {
 void
 psset_alloc_new(psset_t *psset, hpdata_t *ps, edata_t *r_edata, size_t size) {
 	hpdata_assert_empty(ps);
-	size_t npages = size / PAGE;
-	size_t begin = hpdata_reserve_alloc(ps, npages);
-	uintptr_t addr = (uintptr_t)hpdata_addr_get(ps) + begin * PAGE;
-	edata_init(r_edata, edata_arena_ind_get(r_edata), (void *)addr, size,
+	void *addr = hpdata_reserve_alloc(ps, size);
+	edata_init(r_edata, edata_arena_ind_get(r_edata), addr, size,
 	    /* slab */ false, SC_NSIZES, /* sn */ 0, extent_state_active,
 	    /* zeroed */ false, /* committed */ true, EXTENT_PAI_HPA,
 	    EXTENT_NOT_HEAD);
@@ -171,10 +169,9 @@ psset_alloc_reuse(psset_t *psset, edata_t *r_edata, size_t size) {
        }
        psset_remove(psset, ps);
 
-       size_t npages = size / PAGE;
-       size_t begin = hpdata_reserve_alloc(ps, npages);
-       uintptr_t addr = (uintptr_t)hpdata_addr_get(ps) + begin * PAGE;
-       edata_init(r_edata, edata_arena_ind_get(r_edata), (void *)addr, size,
+
+       void *addr = hpdata_reserve_alloc(ps, size);
+       edata_init(r_edata, edata_arena_ind_get(r_edata), addr, size,
 	   /* slab */ false, SC_NSIZES, /* sn */ 0, extent_state_active,
 	   /* zeroed */ false, /* committed */ true, EXTENT_PAI_HPA,
 	   EXTENT_NOT_HEAD);
@@ -190,13 +187,8 @@ psset_dalloc(psset_t *psset, edata_t *edata) {
 	assert(edata_ps_get(edata) != NULL);
 	hpdata_t *ps = edata_ps_get(edata);
 
-	size_t begin =
-	    ((uintptr_t)edata_base_get(edata) - (uintptr_t)hpdata_addr_get(ps))
-	    >> LG_PAGE;
-	size_t len = edata_size_get(edata) >> LG_PAGE;
-
 	psset_remove(psset, ps);
-	hpdata_unreserve(ps, begin, len);
+	hpdata_unreserve(ps, edata_base_get(edata), edata_size_get(edata));
 	if (hpdata_empty(ps)) {
 		return ps;
 	} else {