lib/memtrace/src/mappings.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 
  3 pub const Removal = struct {
  4     requested_bytes: usize,
  5     tracked_bytes: usize,
  6 
  7     pub fn untrackedBytes(self: Removal) usize {
  8         return self.requested_bytes - self.tracked_bytes;
  9     }
 10 };
 11 
 12 const Mapping = struct {
 13     start: usize,
 14     end: usize,
 15 };
 16 
 17 pub const Ledger = struct {
 18     ranges: std.ArrayListUnmanaged(Mapping) = .empty,
 19     mapped_bytes: usize = 0,
 20 
 21     pub fn deinit(self: *Ledger, allocator: std.mem.Allocator) void {
 22         self.ranges.deinit(allocator);
 23         self.* = undefined;
 24     }
 25 
 26     pub fn map(
 27         self: *Ledger,
 28         allocator: std.mem.Allocator,
 29         address: usize,
 30         len: usize,
 31     ) !usize {
 32         if (len == 0) return 0;
 33         const end = std.math.add(usize, address, len) catch
 34             return error.InvalidMappingRange;
 35         try self.ranges.ensureUnusedCapacity(allocator, 2);
 36         const displaced = self.removeRange(address, end);
 37         var start = address;
 38         var merged_end = end;
 39         var index = self.insertionIndex(address);
 40         if (index > 0 and self.ranges.items[index - 1].end == start) {
 41             start = self.ranges.items[index - 1].start;
 42             _ = self.ranges.orderedRemove(index - 1);
 43             index -= 1;
 44         }
 45         if (index < self.ranges.items.len and
 46             self.ranges.items[index].start == merged_end)
 47         {
 48             merged_end = self.ranges.items[index].end;
 49             _ = self.ranges.orderedRemove(index);
 50         }
 51         self.ranges.insertAssumeCapacity(index, .{
 52             .start = start,
 53             .end = merged_end,
 54         });
 55         self.mapped_bytes = self.mapped_bytes - displaced + len;
 56         return displaced;
 57     }
 58 
 59     pub fn unmap(
 60         self: *Ledger,
 61         allocator: std.mem.Allocator,
 62         address: usize,
 63         len: usize,
 64     ) !Removal {
 65         if (len == 0) return .{
 66             .requested_bytes = 0,
 67             .tracked_bytes = 0,
 68         };
 69         const end = std.math.add(usize, address, len) catch
 70             return error.InvalidMappingRange;
 71         try self.ranges.ensureUnusedCapacity(allocator, 1);
 72         const tracked = self.removeRange(address, end);
 73         self.mapped_bytes -= tracked;
 74         return .{
 75             .requested_bytes = len,
 76             .tracked_bytes = tracked,
 77         };
 78     }
 79 
 80     pub fn count(self: *const Ledger) usize {
 81         return self.ranges.items.len;
 82     }
 83 
 84     fn removeRange(self: *Ledger, start: usize, end: usize) usize {
 85         var removed: usize = 0;
 86         var index: usize = 0;
 87         while (index < self.ranges.items.len) {
 88             const current = self.ranges.items[index];
 89             if (current.end <= start) {
 90                 index += 1;
 91                 continue;
 92             }
 93             if (current.start >= end) break;
 94             const overlap_start = @max(current.start, start);
 95             const overlap_end = @min(current.end, end);
 96             removed += overlap_end - overlap_start;
 97             if (start <= current.start and end >= current.end) {
 98                 _ = self.ranges.orderedRemove(index);
 99                 continue;
100             }
101             if (start <= current.start) {
102                 self.ranges.items[index].start = end;
103                 index += 1;
104                 continue;
105             }
106             if (end >= current.end) {
107                 self.ranges.items[index].end = start;
108                 index += 1;
109                 continue;
110             }
111             self.ranges.items[index].end = start;
112             self.ranges.insertAssumeCapacity(index + 1, .{
113                 .start = end,
114                 .end = current.end,
115             });
116             index += 2;
117         }
118         return removed;
119     }
120 
121     fn insertionIndex(self: *const Ledger, address: usize) usize {
122         var low: usize = 0;
123         var high = self.ranges.items.len;
124         while (low < high) {
125             const middle = low + (high - low) / 2;
126             if (self.ranges.items[middle].start < address) {
127                 low = middle + 1;
128             } else {
129                 high = middle;
130             }
131         }
132         return low;
133     }
134 };
135 
136 test "mapping ledger handles trim split replacement and adjacency" {
137     var ledger: Ledger = .{};
138     defer ledger.deinit(std.testing.allocator);
139 
140     try std.testing.expectEqual(
141         @as(usize, 0),
142         try ledger.map(std.testing.allocator, 0x1000, 0x4000),
143     );
144     const prefix = try ledger.unmap(
145         std.testing.allocator,
146         0x1000,
147         0x1000,
148     );
149     try std.testing.expectEqual(@as(usize, 0), prefix.untrackedBytes());
150     const middle = try ledger.unmap(
151         std.testing.allocator,
152         0x3000,
153         0x1000,
154     );
155     try std.testing.expectEqual(@as(usize, 0x1000), middle.tracked_bytes);
156     try std.testing.expectEqual(@as(usize, 0x2000), ledger.mapped_bytes);
157     try std.testing.expectEqual(@as(usize, 2), ledger.count());
158 
159     try std.testing.expectEqual(
160         @as(usize, 0x1000),
161         try ledger.map(std.testing.allocator, 0x3000, 0x2000),
162     );
163     try std.testing.expectEqual(@as(usize, 0x3000), ledger.mapped_bytes);
164     try std.testing.expectEqual(@as(usize, 1), ledger.count());
165 }