Skip to content

Commit 12ed288

Browse files
yordisyordis
committed
Store map tuple layout in TypeMap
Compute map entry ABI and value offsets once during type building, and reuse that metadata in runtime map lift/lower paths instead of recalculating tuple layout at each call site.
1 parent 3c87977 commit 12ed288

4 files changed

Lines changed: 187 additions & 79 deletions

File tree

‎crates/environ/src/component/types.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1195,6 +1195,14 @@ pub struct TypeMap {
11951195
pub key: InterfaceType,
11961196
/// The value type of the map.
11971197
pub value: InterfaceType,
1198+
/// Byte information for each map entry represented as `tuple<key, value>`.
1199+
pub entry_abi: CanonicalAbiInfo,
1200+
/// Offset in bytes from the start of the entry tuple to the value field in
1201+
/// memory32.
1202+
pub value_offset32: u32,
1203+
/// Offset in bytes from the start of the entry tuple to the value field in
1204+
/// memory64.
1205+
pub value_offset64: u32,
11981206
}
11991207

12001208
/// Shape of a "fixed size list" interface type.

‎crates/environ/src/component/types_builder.rs‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -701,9 +701,24 @@ impl ComponentTypesBuilder {
701701
assert_eq!(types.id(), self.module_types.validator_id());
702702
let key_ty = self.valtype(types, key)?;
703703
let value_ty = self.valtype(types, value)?;
704+
let key_abi = self.component_types.canonical_abi(&key_ty);
705+
let value_abi = self.component_types.canonical_abi(&value_ty);
706+
let entry_abi = CanonicalAbiInfo::record([key_abi, value_abi].into_iter());
707+
708+
let mut offset32 = 0;
709+
key_abi.next_field32(&mut offset32);
710+
let value_offset32 = value_abi.next_field32(&mut offset32);
711+
712+
let mut offset64 = 0;
713+
key_abi.next_field64(&mut offset64);
714+
let value_offset64 = value_abi.next_field64(&mut offset64);
715+
704716
Ok(self.add_map_type(TypeMap {
705717
key: key_ty,
706718
value: value_ty,
719+
entry_abi,
720+
value_offset32,
721+
value_offset64,
707722
}))
708723
}
709724

‎crates/wasmtime/src/runtime/component/func/typed.rs‎

Lines changed: 151 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -2155,32 +2155,27 @@ fn lower_map_iter<'a, K, V, U>(
21552155
cx: &mut LowerContext<'_, U>,
21562156
key_ty: InterfaceType,
21572157
value_ty: InterfaceType,
2158+
tuple_size: usize,
2159+
tuple_align: u32,
2160+
value_offset: usize,
21582161
len: usize,
21592162
iter: impl Iterator<Item = (&'a K, &'a V)>,
21602163
) -> Result<(usize, usize)>
21612164
where
21622165
K: Lower + 'a,
21632166
V: Lower + 'a,
21642167
{
2165-
// Calculate the tuple layout: each entry is a (key, value) record.
2166-
let tuple_abi = CanonicalAbiInfo::record_static(&[K::ABI, V::ABI]);
2167-
let tuple_size = tuple_abi.size32 as usize;
2168-
let tuple_align = tuple_abi.align32;
2169-
21702168
let size = len
21712169
.checked_mul(tuple_size)
21722170
.ok_or_else(|| format_err!("size overflow copying a map"))?;
21732171
let ptr = cx.realloc(0, 0, tuple_align, size)?;
21742172

21752173
let mut entry_offset = ptr;
21762174
for (key, value) in iter {
2177-
// Lower key at the start of the tuple
2178-
let mut field_offset = 0usize;
2179-
let key_field = K::ABI.next_field32_size(&mut field_offset);
2180-
<K as Lower>::linear_lower_to_memory(key, cx, key_ty, entry_offset + key_field)?;
2181-
// Lower value at its aligned offset within the tuple
2182-
let value_field = V::ABI.next_field32_size(&mut field_offset);
2183-
<V as Lower>::linear_lower_to_memory(value, cx, value_ty, entry_offset + value_field)?;
2175+
// Keys are the first field in each entry tuple.
2176+
<K as Lower>::linear_lower_to_memory(key, cx, key_ty, entry_offset)?;
2177+
// Values start at the precomputed value offset within the tuple.
2178+
<V as Lower>::linear_lower_to_memory(value, cx, value_ty, entry_offset + value_offset)?;
21842179
entry_offset += tuple_size;
21852180
}
21862181

@@ -2198,14 +2193,29 @@ where
21982193
K: Lower + 'a,
21992194
V: Lower + 'a,
22002195
{
2201-
let (key_ty, value_ty) = match ty {
2196+
let (key_ty, value_ty, tuple_size, tuple_align, value_offset) = match ty {
22022197
InterfaceType::Map(i) => {
22032198
let m = &cx.types[i];
2204-
(m.key, m.value)
2199+
(
2200+
m.key,
2201+
m.value,
2202+
usize::try_from(m.entry_abi.size32).unwrap(),
2203+
m.entry_abi.align32,
2204+
usize::try_from(m.value_offset32).unwrap(),
2205+
)
22052206
}
22062207
_ => bad_type_info(),
22072208
};
2208-
let (ptr, len) = lower_map_iter(cx, key_ty, value_ty, len, iter)?;
2209+
let (ptr, len) = lower_map_iter(
2210+
cx,
2211+
key_ty,
2212+
value_ty,
2213+
tuple_size,
2214+
tuple_align,
2215+
value_offset,
2216+
len,
2217+
iter,
2218+
)?;
22092219
// See "WRITEPTR64" above for why this is always storing a 64-bit integer.
22102220
map_maybe_uninit!(dst[0]).write(ValRaw::i64(ptr as i64));
22112221
map_maybe_uninit!(dst[1]).write(ValRaw::i64(len as i64));
@@ -2223,15 +2233,30 @@ where
22232233
K: Lower + 'a,
22242234
V: Lower + 'a,
22252235
{
2226-
let (key_ty, value_ty) = match ty {
2236+
let (key_ty, value_ty, tuple_size, tuple_align, value_offset) = match ty {
22272237
InterfaceType::Map(i) => {
22282238
let m = &cx.types[i];
2229-
(m.key, m.value)
2239+
(
2240+
m.key,
2241+
m.value,
2242+
usize::try_from(m.entry_abi.size32).unwrap(),
2243+
m.entry_abi.align32,
2244+
usize::try_from(m.value_offset32).unwrap(),
2245+
)
22302246
}
22312247
_ => bad_type_info(),
22322248
};
22332249
debug_assert!(offset % (CanonicalAbiInfo::POINTER_PAIR.align32 as usize) == 0);
2234-
let (ptr, len) = lower_map_iter(cx, key_ty, value_ty, len, iter)?;
2250+
let (ptr, len) = lower_map_iter(
2251+
cx,
2252+
key_ty,
2253+
value_ty,
2254+
tuple_size,
2255+
tuple_align,
2256+
value_offset,
2257+
len,
2258+
iter,
2259+
)?;
22352260
*cx.get(offset + 0) = u32::try_from(ptr).unwrap().to_le_bytes();
22362261
*cx.get(offset + 4) = u32::try_from(len).unwrap().to_le_bytes();
22372262
Ok(())
@@ -2248,29 +2273,50 @@ where
22482273
ty: InterfaceType,
22492274
src: &Self::Lower,
22502275
) -> Result<Self> {
2251-
let (key_ty, value_ty) = match ty {
2276+
let (key_ty, value_ty, tuple_size, tuple_align, value_offset) = match ty {
22522277
InterfaceType::Map(i) => {
22532278
let m = &cx.types[i];
2254-
(m.key, m.value)
2279+
(
2280+
m.key,
2281+
m.value,
2282+
usize::try_from(m.entry_abi.size32).unwrap(),
2283+
usize::try_from(m.entry_abi.align32).unwrap(),
2284+
usize::try_from(m.value_offset32).unwrap(),
2285+
)
22552286
}
22562287
_ => bad_type_info(),
22572288
};
22582289
// FIXME(#4311): needs memory64 treatment
22592290
let ptr = src[0].get_u32();
22602291
let len = src[1].get_u32();
22612292
let (ptr, len) = (usize::try_from(ptr)?, usize::try_from(len)?);
2262-
lift_map(cx, key_ty, value_ty, ptr, len)
2293+
lift_map(
2294+
cx,
2295+
key_ty,
2296+
value_ty,
2297+
tuple_size,
2298+
tuple_align,
2299+
value_offset,
2300+
ptr,
2301+
len,
2302+
)
22632303
}
22642304

22652305
fn linear_lift_from_memory(
22662306
cx: &mut LiftContext<'_>,
22672307
ty: InterfaceType,
22682308
bytes: &[u8],
22692309
) -> Result<Self> {
2270-
let (key_ty, value_ty) = match ty {
2310+
let (key_ty, value_ty, tuple_size, tuple_align, value_offset) = match ty {
22712311
InterfaceType::Map(i) => {
22722312
let m = &cx.types[i];
2273-
(m.key, m.value)
2313+
(
2314+
m.key,
2315+
m.value,
2316+
usize::try_from(m.entry_abi.size32).unwrap(),
2317+
usize::try_from(m.entry_abi.align32).unwrap(),
2318+
usize::try_from(m.value_offset32).unwrap(),
2319+
)
22742320
}
22752321
_ => bad_type_info(),
22762322
};
@@ -2279,7 +2325,16 @@ where
22792325
let ptr = u32::from_le_bytes(bytes[..4].try_into().unwrap());
22802326
let len = u32::from_le_bytes(bytes[4..].try_into().unwrap());
22812327
let (ptr, len) = (usize::try_from(ptr)?, usize::try_from(len)?);
2282-
lift_map(cx, key_ty, value_ty, ptr, len)
2328+
lift_map(
2329+
cx,
2330+
key_ty,
2331+
value_ty,
2332+
tuple_size,
2333+
tuple_align,
2334+
value_offset,
2335+
ptr,
2336+
len,
2337+
)
22832338
}
22842339
}
22852340

@@ -2289,6 +2344,9 @@ fn lift_map_pairs<K, V>(
22892344
cx: &mut LiftContext<'_>,
22902345
key_ty: InterfaceType,
22912346
value_ty: InterfaceType,
2347+
tuple_size: usize,
2348+
tuple_align: usize,
2349+
value_offset: usize,
22922350
ptr: usize,
22932351
len: usize,
22942352
mut insert: impl FnMut(K, V) -> Result<()>,
@@ -2297,10 +2355,6 @@ where
22972355
K: Lift,
22982356
V: Lift,
22992357
{
2300-
let tuple_abi = CanonicalAbiInfo::record_static(&[K::ABI, V::ABI]);
2301-
let tuple_size = tuple_abi.size32 as usize;
2302-
let tuple_align = tuple_abi.align32 as usize;
2303-
23042358
match len
23052359
.checked_mul(tuple_size)
23062360
.and_then(|total| ptr.checked_add(total))
@@ -2315,13 +2369,10 @@ where
23152369
for i in 0..len {
23162370
let entry_base = ptr + (i * tuple_size);
23172371

2318-
let mut field_offset = 0usize;
2319-
let key_field = K::ABI.next_field32_size(&mut field_offset);
2320-
let key_bytes = &cx.memory()[entry_base + key_field..][..K::SIZE32];
2372+
let key_bytes = &cx.memory()[entry_base..][..K::SIZE32];
23212373
let key = K::linear_lift_from_memory(cx, key_ty, key_bytes)?;
23222374

2323-
let value_field = V::ABI.next_field32_size(&mut field_offset);
2324-
let value_bytes = &cx.memory()[entry_base + value_field..][..V::SIZE32];
2375+
let value_bytes = &cx.memory()[entry_base + value_offset..][..V::SIZE32];
23252376
let value = V::linear_lift_from_memory(cx, value_ty, value_bytes)?;
23262377

23272378
insert(key, value)?;
@@ -2335,6 +2386,9 @@ fn lift_map<K, V>(
23352386
cx: &mut LiftContext<'_>,
23362387
key_ty: InterfaceType,
23372388
value_ty: InterfaceType,
2389+
tuple_size: usize,
2390+
tuple_align: usize,
2391+
value_offset: usize,
23382392
ptr: usize,
23392393
len: usize,
23402394
) -> Result<HashMap<K, V>>
@@ -2343,10 +2397,20 @@ where
23432397
V: Lift,
23442398
{
23452399
let mut result = HashMap::with_capacity(len);
2346-
lift_map_pairs(cx, key_ty, value_ty, ptr, len, |k, v| {
2347-
result.insert(k, v);
2348-
Ok(())
2349-
})?;
2400+
lift_map_pairs(
2401+
cx,
2402+
key_ty,
2403+
value_ty,
2404+
tuple_size,
2405+
tuple_align,
2406+
value_offset,
2407+
ptr,
2408+
len,
2409+
|k, v| {
2410+
result.insert(k, v);
2411+
Ok(())
2412+
},
2413+
)?;
23502414
Ok(result)
23512415
}
23522416

@@ -2475,29 +2539,50 @@ where
24752539
ty: InterfaceType,
24762540
src: &Self::Lower,
24772541
) -> Result<Self> {
2478-
let (key_ty, value_ty) = match ty {
2542+
let (key_ty, value_ty, tuple_size, tuple_align, value_offset) = match ty {
24792543
InterfaceType::Map(i) => {
24802544
let m = &cx.types[i];
2481-
(m.key, m.value)
2545+
(
2546+
m.key,
2547+
m.value,
2548+
usize::try_from(m.entry_abi.size32).unwrap(),
2549+
usize::try_from(m.entry_abi.align32).unwrap(),
2550+
usize::try_from(m.value_offset32).unwrap(),
2551+
)
24822552
}
24832553
_ => bad_type_info(),
24842554
};
24852555
// FIXME(#4311): needs memory64 treatment
24862556
let ptr = src[0].get_u32();
24872557
let len = src[1].get_u32();
24882558
let (ptr, len) = (usize::try_from(ptr)?, usize::try_from(len)?);
2489-
lift_try_map(cx, key_ty, value_ty, ptr, len)
2559+
lift_try_map(
2560+
cx,
2561+
key_ty,
2562+
value_ty,
2563+
tuple_size,
2564+
tuple_align,
2565+
value_offset,
2566+
ptr,
2567+
len,
2568+
)
24902569
}
24912570

24922571
fn linear_lift_from_memory(
24932572
cx: &mut LiftContext<'_>,
24942573
ty: InterfaceType,
24952574
bytes: &[u8],
24962575
) -> Result<Self> {
2497-
let (key_ty, value_ty) = match ty {
2576+
let (key_ty, value_ty, tuple_size, tuple_align, value_offset) = match ty {
24982577
InterfaceType::Map(i) => {
24992578
let m = &cx.types[i];
2500-
(m.key, m.value)
2579+
(
2580+
m.key,
2581+
m.value,
2582+
usize::try_from(m.entry_abi.size32).unwrap(),
2583+
usize::try_from(m.entry_abi.align32).unwrap(),
2584+
usize::try_from(m.value_offset32).unwrap(),
2585+
)
25012586
}
25022587
_ => bad_type_info(),
25032588
};
@@ -2506,7 +2591,16 @@ where
25062591
let ptr = u32::from_le_bytes(bytes[..4].try_into().unwrap());
25072592
let len = u32::from_le_bytes(bytes[4..].try_into().unwrap());
25082593
let (ptr, len) = (usize::try_from(ptr)?, usize::try_from(len)?);
2509-
lift_try_map(cx, key_ty, value_ty, ptr, len)
2594+
lift_try_map(
2595+
cx,
2596+
key_ty,
2597+
value_ty,
2598+
tuple_size,
2599+
tuple_align,
2600+
value_offset,
2601+
ptr,
2602+
len,
2603+
)
25102604
}
25112605
}
25122606

@@ -2515,6 +2609,9 @@ fn lift_try_map<K, V>(
25152609
cx: &mut LiftContext<'_>,
25162610
key_ty: InterfaceType,
25172611
value_ty: InterfaceType,
2612+
tuple_size: usize,
2613+
tuple_align: usize,
2614+
value_offset: usize,
25182615
ptr: usize,
25192616
len: usize,
25202617
) -> Result<TryHashMap<K, V>>
@@ -2523,9 +2620,17 @@ where
25232620
V: Lift,
25242621
{
25252622
let mut result = TryHashMap::with_capacity(len)?;
2526-
lift_map_pairs(cx, key_ty, value_ty, ptr, len, |k, v| {
2527-
result.insert(k, v).map(drop).map_err(Into::into)
2528-
})?;
2623+
lift_map_pairs(
2624+
cx,
2625+
key_ty,
2626+
value_ty,
2627+
tuple_size,
2628+
tuple_align,
2629+
value_offset,
2630+
ptr,
2631+
len,
2632+
|k, v| result.insert(k, v).map(drop).map_err(Into::into),
2633+
)?;
25292634
Ok(result)
25302635
}
25312636

0 commit comments

Comments
 (0)