complete homework 1

include/core/allocator.h

@@ -26,6 +26,11 @@ namespace infini {
     // =================================== 作业 ===================================
     // TODO：可能需要设计一个数据结构来存储free block，以便于管理和合并
     // HINT: 可以使用一个 map 来存储 free block，key 为 block 的起始/结尾地址，value 为 block 的大小
+    // Free block management: 
+    // - freeBlocksByAddr: map from start address to size, for allocation search
+    // - freeBlocksByEnd: map from end address to size, for merging adjacent blocks
+    std::map<size_t, size_t> freeBlocksByAddr;  // key: start address, value: size
+    std::map<size_t, size_t> freeBlocksByEnd;   // key: end address, value: size
     // =================================== 作业 ===================================
 
   public:

src/core/allocator.cc

@@ -31,6 +31,45 @@ namespace infini
 
         // =================================== 作业 ===================================
         // TODO: 设计一个算法来分配内存，返回起始地址偏移量
+        // Use First Fit algorithm: find the first free block that is large enough
+        for (auto it = freeBlocksByAddr.begin(); it != freeBlocksByAddr.end(); ++it)
+        {
+            size_t addr = it->first;
+            size_t blockSize = it->second;
+
+            if (blockSize >= size)
+            {
+                // Found a suitable block
+                // Remove this block from both maps
+                freeBlocksByAddr.erase(it);
+                freeBlocksByEnd.erase(addr + blockSize);
+
+                // If the block is larger than needed, add the remaining part back
+                if (blockSize > size)
+                {
+                    size_t newAddr = addr + size;
+                    size_t newSize = blockSize - size;
+                    freeBlocksByAddr[newAddr] = newSize;
+                    freeBlocksByEnd[newAddr + newSize] = newSize;
+                }
+
+                // Update memory usage statistics
+                used += size;
+                if (used > peak)
+                {
+                    peak = used;
+                }
+
+                return addr;
+            }
+        }
+
+        // No suitable free block found, allocate at the end
+        size_t addr = peak;
+        used += size;
+        peak += size;
+
+        return addr;
         // =================================== 作业 ===================================
 
         return 0;
@@ -43,6 +82,75 @@ namespace infini
 
         // =================================== 作业 ===================================
         // TODO: 设计一个算法来回收内存
+        // Update memory usage
+        used -= size;
+
+        size_t blockStart = addr;
+        size_t blockEnd = addr + size;
+        size_t blockSize = size;
+
+        // Special case: if freeing the block at the end, just reduce peak
+        if (blockEnd == peak)
+        {
+            // Check if we can merge with a previous free block at the end
+            auto prevIt = freeBlocksByEnd.find(blockStart);
+            if (prevIt != freeBlocksByEnd.end())
+            {
+                // Merge with the previous block and reduce peak further
+                size_t prevSize = prevIt->second;
+                size_t prevStart = blockStart - prevSize;
+
+                // Remove the previous block from both maps
+                freeBlocksByAddr.erase(prevStart);
+                freeBlocksByEnd.erase(blockStart);
+
+                // Reduce peak to the start of the merged block
+                peak = prevStart;
+            }
+            else
+            {
+                // Just reduce peak
+                peak = blockStart;
+            }
+            return;
+        }
+
+        // Try to merge with the previous adjacent free block
+        auto prevIt = freeBlocksByEnd.find(blockStart);
+        if (prevIt != freeBlocksByEnd.end())
+        {
+            // Found a previous adjacent block
+            size_t prevSize = prevIt->second;
+            size_t prevStart = blockStart - prevSize;
+
+            // Remove the previous block from both maps
+            freeBlocksByAddr.erase(prevStart);
+            freeBlocksByEnd.erase(blockStart);
+
+            // Merge: extend the current block backwards
+            blockStart = prevStart;
+            blockSize += prevSize;
+        }
+
+        // Try to merge with the next adjacent free block
+        auto nextIt = freeBlocksByAddr.find(blockEnd);
+        if (nextIt != freeBlocksByAddr.end())
+        {
+            // Found a next adjacent block
+            size_t nextSize = nextIt->second;
+
+            // Remove the next block from both maps
+            freeBlocksByAddr.erase(blockEnd);
+            freeBlocksByEnd.erase(blockEnd + nextSize);
+
+            // Merge: extend the current block forwards
+            blockSize += nextSize;
+            blockEnd += nextSize;
+        }
+
+        // Add the merged (or original) free block to both maps
+        freeBlocksByAddr[blockStart] = blockSize;
+        freeBlocksByEnd[blockEnd] = blockSize;
         // =================================== 作业 ===================================
     }