add get_seg class of functions to sdk_comp

Author: rbran

src/id0/db.rs

@@ -460,6 +460,7 @@ impl<K: IDAKind> ID0Section<K> {
             .map(|x| x.map(|x| SegmentIdx(x.0)))
     }
 
+    /// Return an iterator of segments, ordered by address start
     pub fn segments(&self, idx: SegmentIdx<K>) -> SegmentIter<'_, K> {
         let segments = self.sup_range(idx.into(), ARRAY_SUP_TAG);
         SegmentIter {
@@ -502,17 +503,20 @@ impl<K: IDAKind> ID0Section<K> {
     // TODO there is also a "P" entry in patches, it seems to only contains
     // the value 0x01 for each equivalent "A" entry
 
-    pub fn segment_name(&self, idx: SegmentNameIdx<K>) -> Result<IDBStr<'_>> {
+    pub fn segment_name(
+        &self,
+        idx: SegmentNameIdx<K>,
+    ) -> Result<Option<IDBStr<'_>>> {
         let seg_idx = self.segment_strings_idx()?;
         // TODO I think this is dependent on the version, and not on availability
         if let Some(seg_idx) = seg_idx {
             for seg in self.segment_strings(seg_idx) {
                 let (seg_idx, seg_value) = seg?;
                 if seg_idx.0 == idx.0 {
-                    return Ok(seg_value);
+                    return Ok(Some(seg_value));
                 }
             }
-            Err(anyhow!("Unable to find ID0 Segment Name"))
+            Ok(None)
         } else {
             // if there is no names, AKA `$ segstrings`, search for the key directly
             self.name_by_index(idx)
@@ -522,16 +526,17 @@ impl<K: IDAKind> ID0Section<K> {
     pub(crate) fn name_by_index(
         &self,
         idx: SegmentNameIdx<K>,
-    ) -> Result<IDBStr<'_>> {
+    ) -> Result<Option<IDBStr<'_>>> {
         // if there is no names, AKA `$ segstrings`, search for the key directly
-        let name_idx = self
-            .netnode_tag_idx(
-                NetnodeIdx(K::Usize::from(0xFFu8).swap_bytes() | idx.0),
-                NAME_TAG,
-            )
-            .ok_or_else(|| anyhow!("Not found name for segment {}", idx.0))?;
+        let Some(name_idx) = self.netnode_tag_idx(
+            NetnodeIdx(K::Usize::from(0xFFu8).swap_bytes() | idx.0),
+            NAME_TAG,
+        ) else {
+            return Ok(None);
+        };
         parse_maybe_cstr(&self.entries[name_idx].value)
             .map(IDBStr::new)
+            .map(Option::Some)
             .ok_or_else(|| anyhow!("Invalid segment name {}", idx.0))
     }
 

src/sdk_comp.rs

@@ -1,6 +1,8 @@
 pub mod bytes;
 pub mod frame;
 pub mod funcs;
+/// Functions that are exported on the lib but not part of the SDK
+pub mod inner;
 pub mod lines;
 pub mod nalt;
 pub mod netnode;
@@ -16,3 +18,17 @@ pub enum DataFetch<T> {
     Qty(u16),
     Data(Vec<T>),
 }
+
+pub mod prelude {
+    pub use super::bytes::*;
+    pub use super::frame::*;
+    pub use super::funcs::*;
+    pub use super::inner::*;
+    pub use super::lines::*;
+    pub use super::nalt::*;
+    pub use super::netnode::*;
+    pub use super::pro::*;
+    pub use super::range::*;
+    pub use super::segment::*;
+    pub use super::typeinf::*;
+}

src/sdk_comp/inner.rs

@@ -0,0 +1,25 @@
+use anyhow::Result;
+
+use crate::id0::ID0Section;
+use crate::sdk_comp::segment::segment_t;
+use crate::{IDAKind, IDBStr};
+
+// TODO implement based on the InnerRef
+pub fn get_segm_name<'a, K: IDAKind>(
+    id0: &'a ID0Section<K>,
+    s: &segment_t<K>,
+    flags: i32,
+) -> Result<Option<IDBStr<'a>>> {
+    if flags != 0 {
+        todo!();
+    }
+    id0.segment_name(s.name)
+}
+
+// TODO implement based on the InnerRef
+pub fn get_segm_class<'a, K: IDAKind>(
+    id0: &'a ID0Section<K>,
+    s: &segment_t<K>,
+) -> Result<Option<IDBStr<'a>>> {
+    id0.segment_name(s.sclass)
+}

src/sdk_comp/lines.rs

@@ -16,9 +16,10 @@ pub fn get_sourcefile<K: IDAKind>(
     id0: &ID0Section<K>,
     ea: ea_t<K>,
 ) -> Result<Option<SourceFile<'_, K>>> {
-    let seg = super::segment::getseg_inner(id0, ea)?;
+    let seg =
+        super::segment::getseg_inner(id0, |seg| seg.address.contains(&ea))?;
     if let Some(seg) = seg {
-        let name = id0.segment_name(seg.name).unwrap_or(IDBStr::new(&[]));
+        let name = id0.segment_name(seg.name)?.unwrap_or(IDBStr::new(&[]));
         Ok(Some(SourceFile {
             name,
             address: seg.address,

src/sdk_comp/segment.rs

@@ -2,39 +2,61 @@ use anyhow::Result;
 
 use std::ops::Range;
 
-use crate::id0::{ID0Section, Segment};
+use crate::id0::{
+    ID0Section, Segment, SegmentBitness, SegmentNameIdx, SegmentType,
+};
 use crate::IDAKind;
 
 use super::pro::{bgcolor_t, ea_t, sel_t, uval_t};
 
 #[derive(Debug, Clone)]
 pub struct segment_t<K: IDAKind> {
-    pub range: Range<ea_t<K>>,
-    pub name: uval_t<K>,
-    pub sclass: uval_t<K>,
+    pub address: Range<ea_t<K>>,
+    pub name: SegmentNameIdx<K>,
+    pub sclass: SegmentNameIdx<K>,
     pub orgbase: uval_t<K>,
     pub align: u8,
     pub comb: u8,
     pub perm: u8,
-    pub bitness: u8,
+    pub bitness: SegmentBitness,
     pub flags: u16,
     pub sel: sel_t<K>,
     pub defsr: [sel_t<K>; 16usize],
-    pub type_: u8,
+    pub type_: SegmentType,
     pub color: bgcolor_t,
 }
 
+impl<K: IDAKind> From<Segment<K>> for segment_t<K> {
+    fn from(seg: Segment<K>) -> Self {
+        Self {
+            address: seg.address,
+            name: seg.name,
+            sclass: seg.class_id,
+            orgbase: seg.orgbase,
+            align: seg.align.into(),
+            comb: seg.comb.into(),
+            perm: seg.perm.map(|perm| perm.into_raw()).unwrap_or(0),
+            bitness: seg.bitness.into(),
+            flags: seg.flags.into_raw().into(),
+            sel: seg.selector,
+            defsr: seg.defsr,
+            type_: seg.seg_type.into(),
+            color: seg.color,
+        }
+    }
+}
+
 // InnerRef v9.1 fa53bd30-ebf1-4641-80ef-4ddc73db66cd 0x5de220
 pub(crate) fn getseg_inner<K: IDAKind>(
     id0: &ID0Section<K>,
-    ea: ea_t<K>,
+    mut selector: impl FnMut(&Segment<K>) -> bool,
 ) -> Result<Option<Segment<K>>> {
     let Some(seg_idx) = id0.segments_idx()? else {
         return Ok(None);
     };
     for seg in id0.segments(seg_idx) {
         let seg = seg?;
-        if seg.address.contains(&ea) {
+        if selector(&seg) {
             return Ok(Some(seg));
         }
     }
@@ -46,27 +68,53 @@ pub fn getseg<K: IDAKind>(
     id0: &ID0Section<K>,
     ea: ea_t<K>,
 ) -> Result<Option<segment_t<K>>> {
-    getseg_inner(id0, ea).map(|seg| {
-        seg.map(|seg| segment_t {
-            range: seg.address,
-            name: seg.name.0,
-            sclass: seg.class_id.0,
-            orgbase: seg.orgbase,
-            align: seg.align.into(),
-            comb: seg.comb.into(),
-            perm: seg.perm.map(|perm| perm.into_raw()).unwrap_or(0),
-            bitness: seg.bitness.into(),
-            flags: seg.flags.into_raw().into(),
-            sel: seg.selector,
-            defsr: seg.defsr,
-            type_: seg.seg_type.into(),
-            color: seg.color,
-        })
-    })
+    getseg_inner(id0, |seg| seg.address.contains(&ea))
+        .map(|seg| seg.map(|seg| seg.into()))
 }
 
 // InnerRef v9.1 fa53bd30-ebf1-4641-80ef-4ddc73db66cd 0x5e41e0
 pub fn segtype<K: IDAKind>(id0: &ID0Section<K>, ea: ea_t<K>) -> Result<u8> {
-    getseg_inner(id0, ea)
+    getseg_inner(id0, |seg| seg.address.contains(&ea))
         .map(|seg| seg.map(|seg| seg.seg_type.into()).unwrap_or(8))
 }
+
+pub fn get_segm_qty<K: IDAKind>(id0: &ID0Section<K>) -> Result<usize> {
+    let Some(seg_idx) = id0.segments_idx()? else {
+        return Ok(0);
+    };
+    Ok(id0.segments(seg_idx).count())
+}
+
+pub fn get_next_seg<K: IDAKind>(
+    id0: &ID0Section<K>,
+    ea: ea_t<K>,
+) -> Result<Option<segment_t<K>>> {
+    getseg_inner(id0, |seg| seg.address.start > ea)
+        .map(|seg| seg.map(|seg| seg.into()))
+}
+
+pub fn get_prev_seg<K: IDAKind>(
+    id0: &ID0Section<K>,
+    ea: ea_t<K>,
+) -> Result<Option<segment_t<K>>> {
+    getseg_inner(id0, |seg| seg.address.end < ea)
+        .map(|seg| seg.map(|seg| seg.into()))
+}
+
+pub fn get_first_seg<K: IDAKind>(
+    id0: &ID0Section<K>,
+) -> Result<Option<segment_t<K>>> {
+    getseg_inner(id0, |_seg| true).map(|seg| seg.map(|seg| seg.into()))
+}
+
+pub fn get_last_seg<K: IDAKind>(
+    id0: &ID0Section<K>,
+) -> Result<Option<segment_t<K>>> {
+    let Some(seg_idx) = id0.segments_idx()? else {
+        return Ok(None);
+    };
+    id0.segments(seg_idx)
+        .last()
+        .map(|seg| seg.map(|seg| seg.into()))
+        .transpose()
+}