Merge pull request #726 from aya-rs/btf-iter-alloc

btf: avoid multiple vector allocations

aya-obj/src/btf/types.rs

@@ -76,7 +76,7 @@ impl Const {
 
     pub(crate) fn new(btf_type: u32) -> Self {
         let info = (BtfKind::Const as u32) << 24;
-        Const {
+        Self {
             name_offset: 0,
             info,
             btf_type,
@@ -150,7 +150,7 @@ impl Ptr {
 
     pub fn new(name_offset: u32, btf_type: u32) -> Self {
         let info = (BtfKind::Ptr as u32) << 24;
-        Ptr {
+        Self {
             name_offset,
             info,
             btf_type,
@@ -181,7 +181,7 @@ impl Typedef {
 
     pub(crate) fn new(name_offset: u32, btf_type: u32) -> Self {
         let info = (BtfKind::Typedef as u32) << 24;
-        Typedef {
+        Self {
             name_offset,
             info,
             btf_type,
@@ -211,7 +211,7 @@ impl Float {
 
     pub fn new(name_offset: u32, size: u32) -> Self {
         let info = (BtfKind::Float as u32) << 24;
-        Float {
+        Self {
             name_offset,
             info,
             size,
@@ -262,7 +262,7 @@ impl Func {
     pub fn new(name_offset: u32, proto: u32, linkage: FuncLinkage) -> Self {
         let mut info = (BtfKind::Func as u32) << 24;
         info |= (linkage as u32) & 0xFFFF;
-        Func {
+        Self {
             name_offset,
             info,
             btf_type: proto,
@@ -301,7 +301,7 @@ impl TypeTag {
 
     pub fn new(name_offset: u32, btf_type: u32) -> Self {
         let info = (BtfKind::TypeTag as u32) << 24;
-        TypeTag {
+        Self {
             name_offset,
             info,
             btf_type,
@@ -342,12 +342,19 @@ pub struct Int {
 
 impl Int {
     pub(crate) fn to_bytes(&self) -> Vec<u8> {
-        let mut buf = vec![];
-        buf.extend(bytes_of::<u32>(&self.name_offset));
-        buf.extend(bytes_of::<u32>(&self.info));
-        buf.extend(bytes_of::<u32>(&self.size));
-        buf.extend(bytes_of::<u32>(&self.data));
-        buf
+        let Self {
+            name_offset,
+            info,
+            size,
+            data,
+        } = self;
+        [
+            bytes_of::<u32>(name_offset),
+            bytes_of::<u32>(info),
+            bytes_of::<u32>(size),
+            bytes_of::<u32>(data),
+        ]
+        .concat()
     }
 
     pub(crate) fn kind(&self) -> BtfKind {
@@ -363,7 +370,7 @@ impl Int {
         data |= (encoding as u32 & 0x0f) << 24;
         data |= (offset & 0xff) << 16;
         data |= (size * 8) & 0xff;
-        Int {
+        Self {
             name_offset,
             info,
             size,
@@ -404,15 +411,24 @@ pub struct Enum {
 
 impl Enum {
     pub(crate) fn to_bytes(&self) -> Vec<u8> {
-        let mut buf = vec![];
-        buf.extend(bytes_of::<u32>(&self.name_offset));
-        buf.extend(bytes_of::<u32>(&self.info));
-        buf.extend(bytes_of::<u32>(&self.size));
-        for v in &self.variants {
-            buf.extend(bytes_of::<u32>(&v.name_offset));
-            buf.extend(bytes_of::<u32>(&v.value));
-        }
-        buf
+        let Self {
+            name_offset,
+            info,
+            size,
+            variants,
+        } = self;
+        [
+            bytes_of::<u32>(name_offset),
+            bytes_of::<u32>(info),
+            bytes_of::<u32>(size),
+        ]
+        .into_iter()
+        .chain(variants.iter().flat_map(|BtfEnum { name_offset, value }| {
+            [bytes_of::<u32>(name_offset), bytes_of::<u32>(value)]
+        }))
+        .flatten()
+        .copied()
+        .collect()
     }
 
     pub(crate) fn kind(&self) -> BtfKind {
@@ -426,7 +442,7 @@ impl Enum {
     pub fn new(name_offset: u32, variants: Vec<BtfEnum>) -> Self {
         let mut info = (BtfKind::Enum as u32) << 24;
         info |= (variants.len() as u32) & 0xFFFF;
-        Enum {
+        Self {
             name_offset,
             info,
             size: 4,
@@ -458,16 +474,34 @@ pub struct Enum64 {
 
 impl Enum64 {
     pub(crate) fn to_bytes(&self) -> Vec<u8> {
-        let mut buf = vec![];
-        buf.extend(bytes_of::<u32>(&self.name_offset));
-        buf.extend(bytes_of::<u32>(&self.info));
-        buf.extend(bytes_of::<u32>(&self.size));
-        for v in &self.variants {
-            buf.extend(bytes_of::<u32>(&v.name_offset));
-            buf.extend(bytes_of::<u32>(&v.value_high));
-            buf.extend(bytes_of::<u32>(&v.value_low));
-        }
-        buf
+        let Self {
+            name_offset,
+            info,
+            size,
+            variants,
+        } = self;
+        [
+            bytes_of::<u32>(name_offset),
+            bytes_of::<u32>(info),
+            bytes_of::<u32>(size),
+        ]
+        .into_iter()
+        .chain(variants.iter().flat_map(
+            |BtfEnum64 {
+                 name_offset,
+                 value_low,
+                 value_high,
+             }| {
+                [
+                    bytes_of::<u32>(name_offset),
+                    bytes_of::<u32>(value_low),
+                    bytes_of::<u32>(value_high),
+                ]
+            },
+        ))
+        .flatten()
+        .copied()
+        .collect()
     }
 
     pub(crate) fn kind(&self) -> BtfKind {
@@ -502,16 +536,34 @@ pub struct Struct {
 
 impl Struct {
     pub(crate) fn to_bytes(&self) -> Vec<u8> {
-        let mut buf = vec![];
-        buf.extend(bytes_of::<u32>(&self.name_offset));
-        buf.extend(bytes_of::<u32>(&self.info));
-        buf.extend(bytes_of::<u32>(&self.size));
-        for v in &self.members {
-            buf.extend(bytes_of::<u32>(&v.name_offset));
-            buf.extend(bytes_of::<u32>(&v.btf_type));
-            buf.extend(bytes_of::<u32>(&v.offset));
-        }
-        buf
+        let Self {
+            name_offset,
+            info,
+            size,
+            members,
+        } = self;
+        [
+            bytes_of::<u32>(name_offset),
+            bytes_of::<u32>(info),
+            bytes_of::<u32>(size),
+        ]
+        .into_iter()
+        .chain(members.iter().flat_map(
+            |BtfMember {
+                 name_offset,
+                 btf_type,
+                 offset,
+             }| {
+                [
+                    bytes_of::<u32>(name_offset),
+                    bytes_of::<u32>(btf_type),
+                    bytes_of::<u32>(offset),
+                ]
+            },
+        ))
+        .flatten()
+        .copied()
+        .collect()
     }
 
     pub(crate) fn kind(&self) -> BtfKind {
@@ -525,7 +577,7 @@ impl Struct {
     pub(crate) fn new(name_offset: u32, members: Vec<BtfMember>, size: u32) -> Self {
         let mut info = (BtfKind::Struct as u32) << 24;
         info |= (members.len() as u32) & 0xFFFF;
-        Struct {
+        Self {
             name_offset,
             info,
             size,
@@ -563,16 +615,34 @@ pub struct Union {
 
 impl Union {
     pub(crate) fn to_bytes(&self) -> Vec<u8> {
-        let mut buf = vec![];
-        buf.extend(bytes_of::<u32>(&self.name_offset));
-        buf.extend(bytes_of::<u32>(&self.info));
-        buf.extend(bytes_of::<u32>(&self.size));
-        for v in &self.members {
-            buf.extend(bytes_of::<u32>(&v.name_offset));
-            buf.extend(bytes_of::<u32>(&v.btf_type));
-            buf.extend(bytes_of::<u32>(&v.offset));
-        }
-        buf
+        let Self {
+            name_offset,
+            info,
+            size,
+            members,
+        } = self;
+        [
+            bytes_of::<u32>(name_offset),
+            bytes_of::<u32>(info),
+            bytes_of::<u32>(size),
+        ]
+        .into_iter()
+        .chain(members.iter().flat_map(
+            |BtfMember {
+                 name_offset,
+                 btf_type,
+                 offset,
+             }| {
+                [
+                    bytes_of::<u32>(name_offset),
+                    bytes_of::<u32>(btf_type),
+                    bytes_of::<u32>(offset),
+                ]
+            },
+        ))
+        .flatten()
+        .copied()
+        .collect()
     }
 
     pub(crate) fn kind(&self) -> BtfKind {
@@ -609,6 +679,7 @@ pub(crate) struct BtfArray {
     pub(crate) index_type: u32,
     pub(crate) len: u32,
 }
+
 #[repr(C)]
 #[derive(Clone, Debug)]
 pub struct Array {
@@ -620,12 +691,19 @@ pub struct Array {
 
 impl Array {
     pub(crate) fn to_bytes(&self) -> Vec<u8> {
-        let mut buf = vec![];
-        buf.extend(bytes_of::<u32>(&self.name_offset));
-        buf.extend(bytes_of::<u32>(&self.info));
-        buf.extend(bytes_of::<u32>(&self._unused));
-        buf.extend(bytes_of::<BtfArray>(&self.array));
-        buf
+        let Self {
+            name_offset,
+            info,
+            _unused,
+            array,
+        } = self;
+        [
+            bytes_of::<u32>(name_offset),
+            bytes_of::<u32>(info),
+            bytes_of::<u32>(_unused),
+            bytes_of::<BtfArray>(array),
+        ]
+        .concat()
     }
 
     pub(crate) fn kind(&self) -> BtfKind {
@@ -639,7 +717,7 @@ impl Array {
     #[cfg(test)]
     pub(crate) fn new(name_offset: u32, element_type: u32, index_type: u32, len: u32) -> Self {
         let info = (BtfKind::Array as u32) << 24;
-        Array {
+        Self {
             name_offset,
             info,
             _unused: 0,
@@ -670,15 +748,27 @@ pub struct FuncProto {
 
 impl FuncProto {
     pub(crate) fn to_bytes(&self) -> Vec<u8> {
-        let mut buf = vec![];
-        buf.extend(bytes_of::<u32>(&self.name_offset));
-        buf.extend(bytes_of::<u32>(&self.info));
-        buf.extend(bytes_of::<u32>(&self.return_type));
-        for p in &self.params {
-            buf.extend(bytes_of::<u32>(&p.name_offset));
-            buf.extend(bytes_of::<u32>(&p.btf_type));
-        }
-        buf
+        let Self {
+            name_offset,
+            info,
+            return_type,
+            params,
+        } = self;
+        [
+            bytes_of::<u32>(name_offset),
+            bytes_of::<u32>(info),
+            bytes_of::<u32>(return_type),
+        ]
+        .into_iter()
+        .chain(params.iter().flat_map(
+            |BtfParam {
+                 name_offset,
+                 btf_type,
+             }| { [bytes_of::<u32>(name_offset), bytes_of::<u32>(btf_type)] },
+        ))
+        .flatten()
+        .copied()
+        .collect()
     }
 
     pub(crate) fn kind(&self) -> BtfKind {
@@ -692,7 +782,7 @@ impl FuncProto {
     pub fn new(params: Vec<BtfParam>, return_type: u32) -> Self {
         let mut info = (BtfKind::FuncProto as u32) << 24;
         info |= (params.len() as u32) & 0xFFFF;
-        FuncProto {
+        Self {
             name_offset: 0,
             info,
             return_type,
@@ -732,12 +822,19 @@ pub struct Var {
 
 impl Var {
     pub(crate) fn to_bytes(&self) -> Vec<u8> {
-        let mut buf = vec![];
-        buf.extend(bytes_of::<u32>(&self.name_offset));
-        buf.extend(bytes_of::<u32>(&self.info));
-        buf.extend(bytes_of::<u32>(&self.btf_type));
-        buf.extend(bytes_of::<VarLinkage>(&self.linkage));
-        buf
+        let Self {
+            name_offset,
+            info,
+            btf_type,
+            linkage,
+        } = self;
+        [
+            bytes_of::<u32>(name_offset),
+            bytes_of::<u32>(info),
+            bytes_of::<u32>(btf_type),
+            bytes_of::<VarLinkage>(linkage),
+        ]
+        .concat()
     }
 
     pub(crate) fn kind(&self) -> BtfKind {
@@ -750,7 +847,7 @@ impl Var {
 
     pub fn new(name_offset: u32, btf_type: u32, linkage: VarLinkage) -> Self {
         let info = (BtfKind::Var as u32) << 24;
-        Var {
+        Self {
             name_offset,
             info,
             btf_type,
@@ -778,16 +875,34 @@ pub struct DataSec {
 
 impl DataSec {
     pub(crate) fn to_bytes(&self) -> Vec<u8> {
-        let mut buf = vec![];
-        buf.extend(bytes_of::<u32>(&self.name_offset));
-        buf.extend(bytes_of::<u32>(&self.info));
-        buf.extend(bytes_of::<u32>(&self.size));
-        for e in &self.entries {
-            buf.extend(bytes_of::<u32>(&e.btf_type));
-            buf.extend(bytes_of::<u32>(&e.offset));
-            buf.extend(bytes_of::<u32>(&e.size));
-        }
-        buf
+        let Self {
+            name_offset,
+            info,
+            size,
+            entries,
+        } = self;
+        [
+            bytes_of::<u32>(name_offset),
+            bytes_of::<u32>(info),
+            bytes_of::<u32>(size),
+        ]
+        .into_iter()
+        .chain(entries.iter().flat_map(
+            |DataSecEntry {
+                 btf_type,
+                 offset,
+                 size,
+             }| {
+                [
+                    bytes_of::<u32>(btf_type),
+                    bytes_of::<u32>(offset),
+                    bytes_of::<u32>(size),
+                ]
+            },
+        ))
+        .flatten()
+        .copied()
+        .collect()
     }
 
     pub(crate) fn kind(&self) -> BtfKind {
@@ -801,7 +916,7 @@ impl DataSec {
     pub fn new(name_offset: u32, entries: Vec<DataSecEntry>, size: u32) -> Self {
         let mut info = (BtfKind::DataSec as u32) << 24;
         info |= (entries.len() as u32) & 0xFFFF;
-        DataSec {
+        Self {
             name_offset,
             info,
             size,
@@ -821,12 +936,19 @@ pub struct DeclTag {
 
 impl DeclTag {
     pub(crate) fn to_bytes(&self) -> Vec<u8> {
-        let mut buf = vec![];
-        buf.extend(bytes_of::<u32>(&self.name_offset));
-        buf.extend(bytes_of::<u32>(&self.info));
-        buf.extend(bytes_of::<u32>(&self.btf_type));
-        buf.extend(bytes_of::<i32>(&self.component_index));
-        buf
+        let Self {
+            name_offset,
+            info,
+            btf_type,
+            component_index,
+        } = self;
+        [
+            bytes_of::<u32>(name_offset),
+            bytes_of::<u32>(info),
+            bytes_of::<u32>(btf_type),
+            bytes_of::<i32>(component_index),
+        ]
+        .concat()
     }
 
     pub(crate) fn kind(&self) -> BtfKind {
@@ -839,7 +961,7 @@ impl DeclTag {
 
     pub fn new(name_offset: u32, btf_type: u32, component_index: i32) -> Self {
         let info = (BtfKind::DeclTag as u32) << 24;
-        DeclTag {
+        Self {
             name_offset,
             info,
             btf_type,