feat: allocator POC

joetifa2003 · Apr 4, 2024 · bf51e14 · bf51e14
1 parent 3c45a26
commit bf51e14
Show file tree

Hide file tree

Showing 14 changed files with 177 additions and 312 deletions.
diff --git a/allocator/allocator.go b/allocator/allocator.go
@@ -0,0 +1,32 @@
+package allocator
+
+import (
+	"context"
+	"unsafe"
+)
+
+type ctxKey int
+
+const (
+	AllocatorKey ctxKey = iota
+)
+
+type Allocator interface {
+	Alloc(size int) unsafe.Pointer
+	Realloc(ptr unsafe.Pointer, size int) unsafe.Pointer
+	Free(ptr unsafe.Pointer)
+	Destroy()
+}
+
+func WithAllocator(ctx context.Context, allocator Allocator) context.Context {
+	return context.WithValue(ctx, AllocatorKey, allocator)
+}
+
+func GetAllocator(ctx context.Context) Allocator {
+	allocator, ok := ctx.Value(AllocatorKey).(Allocator)
+	if !ok {
+		return CAllocator{}
+	}
+
+	return allocator
+}
diff --git a/allocator/callocator.go b/allocator/callocator.go
@@ -0,0 +1,23 @@
+package allocator
+
+import (
+	"unsafe"
+
+	"github.com/joetifa2003/mm-go/malloc"
+)
+
+type CAllocator struct{}
+
+func (c CAllocator) Alloc(size int) unsafe.Pointer {
+	return malloc.Malloc(size)
+}
+
+func (c CAllocator) Free(ptr unsafe.Pointer) {
+	malloc.Free(ptr)
+}
+
+func (c CAllocator) Realloc(ptr unsafe.Pointer, size int) unsafe.Pointer {
+	return malloc.Realloc(ptr, int(size))
+}
+
+func (c CAllocator) Destroy() {}
diff --git a/hashmap/hashmap.go b/hashmap/hashmap.go
@@ -1,6 +1,8 @@
 package hashmap
 
 import (
+	"context"
+
 	"github.com/mitchellh/hashstructure/v2"
 
 	"github.com/joetifa2003/mm-go"
@@ -19,17 +21,19 @@ type pair[K comparable, V any] struct {
 type Hashmap[K comparable, V any] struct {
 	pairs      *vector.Vector[*linkedlist.LinkedList[pair[K, V]]]
 	totalTaken int
+	ctx        context.Context
 }
 
 // New creates a new Hashmap with key of type K and value of type V
-func New[K comparable, V any]() *Hashmap[K, V] {
-	hm := mm.Alloc[Hashmap[K, V]]()
-	hm.pairs = vector.New[*linkedlist.LinkedList[pair[K, V]]](8)
+func New[K comparable, V any](ctx context.Context) *Hashmap[K, V] {
+	hm := mm.Alloc[Hashmap[K, V]](ctx)
+	hm.pairs = vector.New[*linkedlist.LinkedList[pair[K, V]]](ctx, 8)
+	hm.ctx = ctx
 	return hm
 }
 
 func (hm *Hashmap[K, V]) extend() {
-	newPairs := vector.New[*linkedlist.LinkedList[pair[K, V]]](hm.pairs.Len() * 2)
+	newPairs := vector.New[*linkedlist.LinkedList[pair[K, V]]](hm.ctx, hm.pairs.Len()*2)
 	oldPairs := hm.pairs
 	defer oldPairs.Free()
 
@@ -67,7 +71,7 @@ func (hm *Hashmap[K, V]) Insert(key K, value V) {
 	idx := int(hash % uint64(hm.pairs.Len()))
 	pairs := hm.pairs.At(idx)
 	if pairs == nil {
-		newPairs := linkedlist.New[pair[K, V]]()
+		newPairs := linkedlist.New[pair[K, V]](hm.ctx)
 		hm.pairs.Set(idx, newPairs)
 		pairs = newPairs
 	}
@@ -194,5 +198,5 @@ func (hm *Hashmap[K, V]) Free() {
 		}
 	}
 	hm.pairs.Free()
-	mm.Free(hm)
+	mm.Free(hm.ctx, hm)
 }
diff --git a/hashmap/hashmap_test.go b/hashmap/hashmap_test.go
@@ -1,6 +1,7 @@
 package hashmap_test
 
 import (
+	"context"
 	"runtime"
 	"testing"
 
@@ -11,11 +12,12 @@ import (
 const TIMES = 5000
 
 func BenchmarkHashmap(b *testing.B) {
+	ctx := context.Background()
 	for i := 0; i < b.N; i++ {
-		h := hashmap.New[int, *mmstring.MMString]()
+		h := hashmap.New[int, mmstring.MMString](ctx)
 
 		for i := 0; i < TIMES; i++ {
-			h.Insert(i, mmstring.From("foo bar"))
+			h.Insert(i, mmstring.From(ctx, "foo bar"))
 		}
 
 		h.Free()

diff --git a/linkedlist/linked_list.go b/linkedlist/linked_list.go
@@ -1,6 +1,7 @@
 package linkedlist
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/joetifa2003/mm-go"
@@ -20,25 +21,27 @@ type LinkedList[T any] struct {
 	head   *linkedListNode[T]
 	tail   *linkedListNode[T]
 	length int
+	ctx    context.Context
 }
 
 // New creates a new linked list.
-func New[T any]() *LinkedList[T] {
-	linkedList := mm.Alloc[LinkedList[T]]()
+func New[T any](ctx context.Context) *LinkedList[T] {
+	linkedList := mm.Alloc[LinkedList[T]](ctx)
+	linkedList.ctx = ctx
 
 	return linkedList
 }
 
 func (ll *LinkedList[T]) init(value T) {
-	ll.head = mm.Alloc[linkedListNode[T]]()
+	ll.head = mm.Alloc[linkedListNode[T]](ll.ctx)
 	ll.head.value = value
 	ll.tail = ll.head
 	ll.length++
 }
 
 func (ll *LinkedList[T]) popLast() T {
 	value := ll.tail.value
-	mm.Free(ll.tail)
+	mm.Free(ll.ctx, ll.tail)
 	ll.tail = nil
 	ll.head = nil
 	ll.length--
@@ -53,7 +56,7 @@ func (ll *LinkedList[T]) PushBack(value T) {
 		return
 	}
 
-	newNode := mm.Alloc[linkedListNode[T]]()
+	newNode := mm.Alloc[linkedListNode[T]](ll.ctx)
 	newNode.value = value
 	newNode.prev = ll.tail
 	ll.tail.next = newNode
@@ -69,7 +72,7 @@ func (ll *LinkedList[T]) PushFront(value T) {
 		return
 	}
 
-	newNode := mm.Alloc[linkedListNode[T]]()
+	newNode := mm.Alloc[linkedListNode[T]](ll.ctx)
 	newNode.value = value
 	newNode.next = ll.head
 	ll.head.prev = newNode
@@ -90,7 +93,7 @@ func (ll *LinkedList[T]) PopBack() T {
 	value := ll.tail.value
 	newTail := ll.tail.prev
 	newTail.next = nil
-	mm.Free(ll.tail)
+	mm.Free(ll.ctx, ll.tail)
 	ll.tail = newTail
 	ll.length--
 
@@ -110,7 +113,7 @@ func (ll *LinkedList[T]) PopFront() T {
 	value := ll.head.value
 	newHead := ll.head.next
 	newHead.prev = nil
-	mm.Free(ll.head)
+	mm.Free(ll.ctx, ll.head)
 	ll.head = newHead
 	ll.length--
 
@@ -172,7 +175,7 @@ func (ll *LinkedList[T]) RemoveAt(idx int) T {
 	prevNode.next = nextNode
 	ll.length--
 
-	mm.Free(node)
+	mm.Free(ll.ctx, node)
 
 	return value
 }
@@ -267,9 +270,9 @@ func (ll *LinkedList[T]) Free() {
 
 	for currentNode != nil {
 		nextNode := currentNode.next
-		mm.Free(currentNode)
+		mm.Free(ll.ctx, currentNode)
 		currentNode = nextNode
 	}
 
-	mm.Free(ll)
+	mm.Free(ll.ctx, ll)
 }
diff --git a/linkedlist/linked_list_test.go b/linkedlist/linked_list_test.go
@@ -1,16 +1,19 @@
 package linkedlist_test
 
 import (
+	"context"
 	"testing"
 
-	"github.com/joetifa2003/mm-go/linkedlist"
 	"github.com/stretchr/testify/assert"
+
+	"github.com/joetifa2003/mm-go/linkedlist"
 )
 
 func testPushAndPop(t *testing.T) {
 	assert := assert.New(t)
+	ctx := context.Background()
 
-	ll := linkedlist.New[int]()
+	ll := linkedlist.New[int](ctx)
 	defer ll.Free()
 
 	ll.PushBack(15)
@@ -47,8 +50,9 @@ func testPushAndPop(t *testing.T) {
 
 func testForEach(t *testing.T) {
 	assert := assert.New(t)
+	ctx := context.Background()
 
-	ll := linkedlist.New[int]()
+	ll := linkedlist.New[int](ctx)
 	defer ll.Free()
 
 	ll.PushBack(2)
@@ -71,8 +75,9 @@ func testForEach(t *testing.T) {
 
 func testIndexing(t *testing.T) {
 	assert := assert.New(t)
+	ctx := context.Background()
 
-	ll := linkedlist.New[int]()
+	ll := linkedlist.New[int](ctx)
 	defer ll.Free()
 
 	ll.PushBack(1)
@@ -110,8 +115,9 @@ func testIndexing(t *testing.T) {
 
 func testRemove(t *testing.T) {
 	assert := assert.New(t)
+	ctx := context.Background()
 
-	ll := linkedlist.New[int]()
+	ll := linkedlist.New[int](ctx)
 	defer ll.Free()
 
 	ll.PushBack(1)

diff --git a/mm.go b/mm.go
@@ -1,9 +1,10 @@
 package mm
 
 import (
+	"context"
 	"unsafe"
 
-	"github.com/joetifa2003/mm-go/malloc"
+	"github.com/joetifa2003/mm-go/allocator"
 )
 
 func getSize[T any]() int {
@@ -12,22 +13,25 @@ func getSize[T any]() int {
 }
 
 // Alloc allocates T and returns a pointer to it.
-func Alloc[T any]() *T {
-	ptr := malloc.Malloc(getSize[T]())
+func Alloc[T any](ctx context.Context) *T {
+	allocator := allocator.GetAllocator(ctx)
+	ptr := allocator.Alloc(getSize[T]())
 	return (*T)(unsafe.Pointer(ptr))
 }
 
 // FreeMany frees memory allocated by Alloc takes a ptr
 // CAUTION: be careful not to double free, and prefer using defer to deallocate
-func Free[T any](ptr *T) {
-	malloc.Free(unsafe.Pointer(ptr))
+func Free[T any](ctx context.Context, ptr *T) {
+	allocator := allocator.GetAllocator(ctx)
+	allocator.Free(unsafe.Pointer(ptr))
 }
 
 // AllocMany allocates n of T and returns a slice representing the heap.
 // CAUTION: don't append to the slice, the purpose of it is to replace pointer
 // arithmetic with slice indexing
-func AllocMany[T any](n int) []T {
-	ptr := malloc.Malloc(getSize[T]() * n)
+func AllocMany[T any](ctx context.Context, n int) []T {
+	allocator := allocator.GetAllocator(ctx)
+	ptr := allocator.Alloc(getSize[T]() * n)
 	return unsafe.Slice(
 		(*T)(ptr),
 		n,
@@ -36,13 +40,15 @@ func AllocMany[T any](n int) []T {
 
 // FreeMany frees memory allocated by AllocMany takes in the slice (aka the heap)
 // CAUTION: be careful not to double free, and prefer using defer to deallocate
-func FreeMany[T any](slice []T) {
-	malloc.Free(unsafe.Pointer(&slice[0]))
+func FreeMany[T any](ctx context.Context, slice []T) {
+	allocator := allocator.GetAllocator(ctx)
+	allocator.Free(unsafe.Pointer(&slice[0]))
 }
 
 // Reallocate reallocates memory allocated with AllocMany and doesn't change underling data
-func Reallocate[T any](slice []T, newN int) []T {
-	ptr := malloc.Realloc(unsafe.Pointer(&slice[0]), getSize[T]()*newN)
+func Reallocate[T any](ctx context.Context, slice []T, newN int) []T {
+	allocator := allocator.GetAllocator(ctx)
+	ptr := allocator.Realloc(unsafe.Pointer(&slice[0]), getSize[T]()*newN)
 	return unsafe.Slice(
 		(*T)(ptr),
 		newN,