# UniversalisOS Memory Allocator Implementation Plan
## Based on PikeOS Source Code Analysis

**Date:** 2026-07-12  
**Source:** pikeos-mirror/sources/ukernel-arm_v7hf/include/mm*.h  
**Target:** kernel/src/core/mm.cpp

---

## PikeOS Memory Management Architecture

### 1. Three-Layer Memory Management

PikeOS uses a three-layer memory management architecture:

```
┌─────────────────────────────────────────────────────────────┐
│  Layer 3: KMEM Allocator (mm_kmem)                         │
│  - Per-partition kernel memory                             │
│  - Used by KDEV drivers and kernel respart page pool       │
└─────────────────────────────────────────────────────────────┘
┌─────────────────────────────────────────────────────────────┐
│  Layer 2: Runtime Allocator (mm_ralloc)                    │
│  - Global and per-partition memory stores                  │
│  - Used for runtime allocations                            │
└─────────────────────────────────────────────────────────────┘
┌─────────────────────────────────────────────────────────────┐
│  Layer 1: Boot Allocator (mm_balloc)                       │
│  - Early boot memory allocation                            │
│  - First-fit strategy                                      │
│  - Cache-line aligned                                      │
└─────────────────────────────────────────────────────────────┘
┌─────────────────────────────────────────────────────────────┐
│  Layer 0: Memory List (mm_list)                            │
│  - Low-level free memory block management                  │
│  - Linked list of free blocks                              │
│  - Merge adjacent blocks                                   │
└─────────────────────────────────────────────────────────────┘
```

---

## 2. Key Data Structures

### P4_mem_store_t (Memory Store)
```c
typedef struct P4_mem_store_str {
    P4_mm_list_t free_list;      // Free memory blocks
    P4_uint32_t id;              // Store ID (partition + index)
    P4_mem_type_t type;          // Privileged / non-privileged
    P4_size_t total;             // Total bytes
    P4_size_t free;              // Free bytes
    P4_spin_t mem_store_lock;    // Lock for runtime allocation
} P4_mem_store_t;
```

### P4_mm_list_t (Memory List)
```c
typedef struct P4_mm_list_str {
    adt_list_t head;             // Linked list of free blocks
} P4_mm_list_t;
```

---

## 3. Key Functions

### Boot Allocator (mm_balloc)
- `mm_balloc_init()` - Initialize boot allocator
- `mm_balloc_assign_mem()` - Assign free memory to boot allocator
- `mm_balloc_assign_tmp()` - Assign temporary memory
- `mm_balloc()` - Allocate memory (panics on failure)
- `mm_balloc_aligned()` - Allocate memory (returns NULL on failure)
- `mm_balloc_phys()` - Allocate by physical address
- `mm_balloc_drain()` - Drain boot allocator
- `mm_balloc_reclaim_tmp()` - Reclaim temporary memory

### Memory List (mm_list)
- `mm_list_init()` - Initialize memory list
- `mm_list_check_overlap()` - Check for overlaps
- `mm_list_assign()` - Add free block to list
- `mm_list_alloc_aligned()` - Allocate with alignment
- `mm_list_alloc_by_addr()` - Allocate by address
- `mm_list_drain()` - Drain memory list

### Runtime Allocator (mm_ralloc)
- `mm_ralloc_boot()` - Allocate from global store at boot
- `mm_ralloc()` - Allocate from memory store

### KMEM Allocator (mm_kmem)
- `mm_kmem_init()` - Initialize KMEM data structures
- `mm_kmem_fill_all()` - Allocate KMEM from stores
- `mm_kmem_alloc()` - Allocate KMEM for partition

---

## 4. Implementation Plan for UniversalisOS

### Phase 1: Memory List (mm_list) - Week 1

**Goal:** Implement low-level free memory block management

**Files to create:**
- `kernel/src/core/mm_list.h`
- `kernel/src/core/mm_list.cpp`

**Key functions:**
```cpp
// Initialize memory list
void mm_list_init(uos_mm_list_t* ml);

// Check overlap
bool mm_list_check_overlap(const uos_mm_list_t* ml, 
                           uos_address_t start, 
                           uos_size_t size);

// Add free block
void mm_list_assign(uos_mm_list_t* ml, 
                    uos_address_t start, 
                    uos_size_t size);

// Allocate with alignment
void* mm_list_alloc_aligned(uos_mm_list_t* ml, 
                            uos_size_t size, 
                            uos_address_t align,
                            uos_address_t destaddr,
                            uos_address_t align_mask);

// Allocate by address
void* mm_list_alloc_by_addr(uos_mm_list_t* ml, 
                            uos_address_t start, 
                            uos_size_t size);

// Drain list
void* mm_list_drain(uos_mm_list_t* ml, uos_size_t* size);
```

**Data structures:**
```cpp
typedef struct uos_mm_list_str {
    uos_list_t head;             // Linked list of free blocks
} uos_mm_list_t;

typedef struct uos_mm_block_str {
    uos_list_node_t node;        // List node
    uos_address_t start;         // Start address
    uos_size_t size;             // Block size
} uos_mm_block_t;
```

---

### Phase 2: Boot Allocator (mm_balloc) - Week 1-2

**Goal:** Implement early boot memory allocation

**Files to create:**
- `kernel/src/core/mm_balloc.h`
- `kernel/src/core/mm_balloc.cpp`

**Key functions:**
```cpp
// Initialize boot allocator
void mm_balloc_init(void);

// Assign free memory
void mm_balloc_assign_mem(uos_phys_addr_t phys_addr, uos_size_t size);

// Assign temporary memory
void mm_balloc_assign_tmp(uos_phys_addr_t phys_addr, uos_size_t size);

// Allocate memory (panics on failure)
void* mm_balloc(uos_size_t size, uos_address_t align);

// Allocate memory (returns NULL on failure)
void* mm_balloc_aligned(uos_size_t size, uos_address_t align);

// Allocate by physical address
void* mm_balloc_phys(uos_phys_addr_t phys_addr, uos_size_t size);

// Drain boot allocator
void* mm_balloc_drain(uos_size_t* size);

// Reclaim temporary memory
void* mm_balloc_reclaim_tmp(uos_size_t* size);
```

**Data structures:**
```cpp
#define BALLOC_NUM_TMP 4

typedef struct uos_balloc_tmp_str {
    uos_address_t start;         // Start address
    uos_size_t size;             // Block size
    bool used;                   // Used flag
} uos_balloc_tmp_t;

static uos_mm_list_t balloc_free_list;
static uos_balloc_tmp_t balloc_tmps[BALLOC_NUM_TMP];
```

---

### Phase 3: Memory Store (mm_store) - Week 2

**Goal:** Implement global and per-partition memory stores

**Files to create:**
- `kernel/src/core/mm_store.h`
- `kernel/src/core/mm_store.cpp`

**Key functions:**
```cpp
// Initialize memory stores
void mm_store_init(void);

// Reclaim temporary memory
void mm_store_reclaim_tmp(void);

// Get store by ID
uos_mem_store_t* mm_store_get_by_id(uos_uint32_t store_id);

// Allocate from store
void* mm_ralloc(uos_mem_store_t* store, 
                uos_size_t size, 
                uos_address_t align,
                uos_address_t destaddr,
                uos_address_t align_mask);

// Allocate from global store at boot
void* mm_ralloc_boot(uos_size_t size, uos_address_t align);
```

**Data structures:**
```cpp
typedef struct uos_mem_store_str {
    uos_mm_list_t free_list;     // Free memory blocks
    uos_uint32_t id;             // Store ID
    uos_mem_type_t type;         // Privileged / non-privileged
    uos_size_t total;            // Total bytes
    uos_size_t free;             // Free bytes
    uos_spin_t lock;             // Lock for runtime allocation
} uos_mem_store_t;

static uos_mem_store_t mm_global_store;
static uos_mem_store_t* mm_part_stores[MAX_PARTITIONS];
```

---

### Phase 4: KMEM Allocator (mm_kmem) - Week 2-3

**Goal:** Implement per-partition kernel memory allocation

**Files to create:**
- `kernel/src/core/mm_kmem.h`
- `kernel/src/core/mm_kmem.cpp`

**Key functions:**
```cpp
// Initialize KMEM data structures
void mm_kmem_init(void);

// Allocate KMEM from stores
void mm_kmem_fill_all(void);

// Allocate KMEM for partition
void* mm_kmem_alloc(uos_uint32_t rp_id, 
                    uos_size_t size, 
                    uos_address_t align);
```

**Data structures:**
```cpp
static uos_mm_list_t* mm_kmem_free_list[MAX_PARTITIONS];
```

---

### Phase 5: Integration with Existing MM - Week 3

**Goal:** Integrate new allocator with existing mm.cpp

**Files to modify:**
- `kernel/src/core/mm.h`
- `kernel/src/core/mm.cpp`

**Key changes:**
1. Replace stub implementations with real allocator calls
2. Add page table management
3. Add COW support
4. Add memory protection

---

## 5. Implementation Details

### Memory List Implementation

```cpp
// mm_list.cpp

#include "mm_list.h"

void mm_list_init(uos_mm_list_t* ml) {
    uos_list_init(&ml->head);
}

bool mm_list_check_overlap(const uos_mm_list_t* ml, 
                           uos_address_t start, 
                           uos_size_t size) {
    uos_mm_block_t* block;
    uos_list_for_each_entry(block, &ml->head, node) {
        if (start < block->start + block->size && 
            start + size > block->start) {
            return true;
        }
    }
    return false;
}

void mm_list_assign(uos_mm_list_t* ml, 
                    uos_address_t start, 
                    uos_size_t size) {
    // Align to cache line
    start = UOS_ALIGN_DOWN(start, UOS_CACHE_LINE_SIZE);
    size = UOS_ALIGN_UP(size, UOS_CACHE_LINE_SIZE);
    
    // Check for overlap
    if (mm_list_check_overlap(ml, start, size)) {
        // Panic or handle error
        return;
    }
    
    // Allocate block structure
    uos_mm_block_t* block = (uos_mm_block_t*)mm_balloc_aligned(
        sizeof(uos_mm_block_t), UOS_CACHE_LINE_SIZE);
    
    block->start = start;
    block->size = size;
    
    // Insert in sorted order
    uos_mm_block_t* pos;
    uos_list_for_each_entry(pos, &ml->head, node) {
        if (pos->start > start) {
            uos_list_add_before(&pos->node, &block->node);
            goto merge;
        }
    }
    uos_list_add_tail(&ml->head, &block->node);
    
merge:
    // Merge adjacent blocks
    uos_mm_block_t* next = uos_list_next_entry(block, node);
    if (next && block->start + block->size == next->start) {
        block->size += next->size;
        uos_list_del(&next->node);
        // Free next block structure
    }
    
    uos_mm_block_t* prev = uos_list_prev_entry(block, node);
    if (prev && prev->start + prev->size == block->start) {
        prev->size += block->size;
        uos_list_del(&block->node);
        // Free block structure
    }
}

void* mm_list_alloc_aligned(uos_mm_list_t* ml, 
                            uos_size_t size, 
                            uos_address_t align,
                            uos_address_t destaddr,
                            uos_address_t align_mask) {
    // Align size to cache line
    size = UOS_ALIGN_UP(size, UOS_CACHE_LINE_SIZE);
    
    uos_mm_block_t* block;
    uos_list_for_each_entry(block, &ml->head, node) {
        // Check if block fits
        uos_address_t aligned_start = UOS_ALIGN_UP(block->start, align);
        uos_size_t aligned_size = block->size - (aligned_start - block->start);
        
        if (aligned_size >= size) {
            // Found suitable block
            if (aligned_size == size) {
                // Exact fit - remove block
                uos_list_del(&block->node);
                return (void*)aligned_start;
            } else {
                // Split block
                block->start = aligned_start + size;
                block->size = aligned_size - size;
                return (void*)aligned_start;
            }
        }
    }
    
    return NULL;  // No suitable block found
}
```

---

## 6. Testing Plan

### Unit Tests
1. **Memory List Tests**
   - Test init
   - Test assign
   - Test overlap check
   - Test alloc_aligned
   - Test alloc_by_addr
   - Test drain
   - Test merge adjacent blocks

2. **Boot Allocator Tests**
   - Test init
   - Test assign_mem
   - Test assign_tmp
   - Test balloc
   - Test balloc_aligned
   - Test balloc_phys
   - Test drain
   - Test reclaim_tmp

3. **Memory Store Tests**
   - Test store_init
   - Test store_get_by_id
   - Test ralloc
   - Test ralloc_boot

4. **KMEM Tests**
   - Test kmem_init
   - Test kmem_fill_all
   - Test kmem_alloc

### Integration Tests
1. **Boot Sequence Test**
   - Initialize boot allocator
   - Assign memory
   - Allocate kernel structures
   - Drain boot allocator
   - Initialize memory stores
   - Initialize KMEM

2. **Runtime Allocation Test**
   - Allocate from global store
   - Allocate from partition store
   - Free memory
   - Check for leaks

---

## 7. Timeline

| Week | Phase | Deliverable |
|------|-------|-------------|
| 1 | Memory List | mm_list.h/cpp with tests |
| 1-2 | Boot Allocator | mm_balloc.h/cpp with tests |
| 2 | Memory Store | mm_store.h/cpp with tests |
| 2-3 | KMEM Allocator | mm_kmem.h/cpp with tests |
| 3 | Integration | Updated mm.h/cpp |
| 3-4 | Testing | All tests passing |

---

## 8. Dependencies

### External Dependencies
- ADT list library (adt/list.h)
- Spinlock library (spinlock.h)
- PSP library (psp.h)

### Internal Dependencies
- kernel/p4types.h
- kernel/p4errorcodes.h
- kernel/p4mm.h

---

## 9. References

- PikeOS mm.h: `pikeos-mirror/sources/ukernel-arm_v7hf/include/mm.h`
- PikeOS mm_balloc.h: `pikeos-mirror/sources/ukernel-arm_v7hf/include/mm_balloc.h`
- PikeOS mm_list.h: `pikeos-mirror/sources/ukernel-arm_v7hf/include/mm_list.h`
- PikeOS mm_kmem.h: `pikeos-mirror/sources/ukernel-arm_v7hf/include/mm_kmem.h`

---

*Plan created: 2026-07-12*
*Based on PikeOS source code analysis*