Add item indexing support to windows-metadata (microsoft#3612)

Author: kennykerr

crates/libs/metadata/readme.md

@@ -19,7 +19,7 @@ Use the Windows metadata support as needed. Here is how you might use the reader
 ```rust,no_run
 use windows_metadata::*;
 
-let index = reader::Index::read("Windows.winmd").unwrap();
+let index = reader::TypeIndex::read("Windows.winmd").unwrap();
 
 let def = index.expect("Windows.Foundation", "Point");
 assert_eq!(def.namespace(), "Windows.Foundation");

crates/libs/metadata/src/reader/blob.rs

@@ -1,7 +1,7 @@
 use super::*;
 
 pub struct Blob<'a> {
-    index: &'a Index,
+    index: &'a TypeIndex,
     file: usize,
     slice: &'a [u8],
 }
@@ -27,7 +27,7 @@ impl std::ops::Deref for Blob<'_> {
 }
 
 impl<'a> Blob<'a> {
-    pub fn new(index: &'a Index, file: usize, slice: &'a [u8]) -> Self {
+    pub fn new(index: &'a TypeIndex, file: usize, slice: &'a [u8]) -> Self {
         Self { index, file, slice }
     }
 

crates/libs/metadata/src/reader/codes.rs

@@ -1,7 +1,7 @@
 use super::*;
 
 pub trait Decode<'a> {
-    fn decode(index: &'a Index, file: usize, code: usize) -> Self;
+    fn decode(index: &'a TypeIndex, file: usize, code: usize) -> Self;
 }
 
 macro_rules! code {
@@ -11,7 +11,7 @@ macro_rules! code {
             $($table($table<'a>),)*
         }
         impl<'a> Decode<'a> for $name<'a> {
-            fn decode(index: &'a Index, file: usize, code: usize) -> Self {
+            fn decode(index: &'a TypeIndex, file: usize, code: usize) -> Self {
                 let (kind, row) = (code & ((1 << $size) - 1), (code >> $size) - 1);
                 match kind {
                     $($code => Self::$table($table(Row::new(index, file, row))),)*

crates/libs/metadata/src/reader/item_index.rs

@@ -0,0 +1,93 @@
+use super::*;
+
+pub enum Item<'a> {
+    Type(TypeDef<'a>),
+    Fn(MethodDef<'a>),
+    Const(Field<'a>),
+}
+
+type HashType<'a> = HashMap<&'a str, HashMap<&'a str, Vec<Item<'a>>>>;
+
+pub struct ItemIndex<'a>(HashType<'a>);
+
+impl<'a> ItemIndex<'a> {
+    pub fn new(index: &'a TypeIndex) -> Self {
+        let mut members: HashType = HashMap::new();
+
+        for (namespace, name, ty) in index.iter() {
+            insert(&mut members, namespace, name, Item::Type(ty));
+
+            if !ty.flags().contains(TypeAttributes::WindowsRuntime) {
+                match ty.category() {
+                    TypeCategory::Class if name == "Apis" => {
+                        for method in ty.methods() {
+                            insert(&mut members, namespace, method.name(), Item::Fn(method));
+                        }
+                        for field in ty.fields() {
+                            insert(&mut members, namespace, field.name(), Item::Const(field));
+                        }
+                    }
+                    TypeCategory::Enum if !ty.has_attribute("ScopedEnumAttribute") => {
+                        for field in ty.fields() {
+                            if field.flags().contains(FieldAttributes::Literal) {
+                                insert(&mut members, namespace, field.name(), Item::Const(field));
+                            }
+                        }
+                    }
+                    _ => {}
+                }
+            }
+        }
+
+        Self(members)
+    }
+
+    pub fn iter(&self) -> impl Iterator<Item = (&str, &str, &Item)> + '_ {
+        self.0
+            .iter()
+            .flat_map(|(namespace, items)| {
+                items
+                    .iter()
+                    .map(move |(name, items)| (namespace, name, items))
+            })
+            .flat_map(|(namespace, name, items)| {
+                items.iter().map(move |item| (*namespace, *name, item))
+            })
+    }
+
+    pub fn items(&self) -> impl Iterator<Item = &Item> + '_ {
+        self.0.values().flat_map(|items| items.values()).flatten()
+    }
+
+    pub fn get(&self, namespace: &str, name: &str) -> impl Iterator<Item = &Item> + '_ {
+        self.0
+            .get(namespace)
+            .and_then(|items| items.get(name))
+            .into_iter()
+            .flatten()
+    }
+
+    #[track_caller]
+    pub fn expect(&self, namespace: &str, name: &str) -> &Item {
+        let mut iter = self.get(namespace, name);
+
+        if let Some(item) = iter.next() {
+            if iter.next().is_none() {
+                item
+            } else {
+                panic!("more than one type found: {namespace}.{name}");
+            }
+        } else {
+            panic!("type not found: {namespace}.{name}")
+        }
+    }
+}
+
+fn insert<'a>(members: &mut HashType<'a>, namespace: &'a str, name: &'a str, member: Item<'a>) {
+    members
+        .entry(namespace)
+        .or_default()
+        .entry(name)
+        .or_default()
+        .push(member);
+}

crates/libs/metadata/src/reader/mod.rs

@@ -3,18 +3,20 @@ use super::*;
 mod blob;
 mod codes;
 mod file;
-mod index;
+mod item_index;
 mod row;
 mod tables;
 mod type_category;
+mod type_index;
 
 pub use blob::*;
 pub use codes::*;
 pub use file::*;
-pub use index::*;
+pub use item_index::*;
 pub use row::*;
 pub use tables::*;
 pub use type_category::*;
+pub use type_index::*;
 
 fn trim_tick(name: &str) -> &str {
     if name.as_bytes().iter().rev().nth(1) == Some(&b'`') {

crates/libs/metadata/src/reader/row.rs

@@ -2,7 +2,7 @@ use super::*;
 
 #[derive(Copy, Clone)]
 pub struct Row<'a> {
-    pub index: &'a Index,
+    pub index: &'a TypeIndex,
     pub file: usize,
     pub pos: usize,
 }
@@ -47,7 +47,7 @@ unsafe impl Send for Row<'_> {}
 unsafe impl Sync for Row<'_> {}
 
 impl<'a> Row<'a> {
-    pub(crate) fn new(index: &'a Index, file: usize, pos: usize) -> Self {
+    pub(crate) fn new(index: &'a TypeIndex, file: usize, pos: usize) -> Self {
         Self { index, file, pos }
     }
 }
@@ -57,7 +57,7 @@ pub trait AsRow<'a>: Copy {
     fn to_row(&self) -> Row<'a>;
     fn from_row(row: Row<'a>) -> Self;
 
-    fn index(&self) -> &'a Index {
+    fn index(&self) -> &'a TypeIndex {
         let row = self.to_row();
         row.index
     }
@@ -129,14 +129,14 @@ pub trait AsRow<'a>: Copy {
 }
 
 pub struct RowIterator<'a, R: AsRow<'a>> {
-    index: &'a Index,
+    index: &'a TypeIndex,
     file: usize,
     rows: std::ops::Range<usize>,
     phantom: std::marker::PhantomData<R>,
 }
 
 impl<'a, R: AsRow<'a>> RowIterator<'a, R> {
-    pub(crate) fn new(index: &'a Index, file: usize, rows: std::ops::Range<usize>) -> Self {
+    pub(crate) fn new(index: &'a TypeIndex, file: usize, rows: std::ops::Range<usize>) -> Self {
         Self {
             index,
             file,

crates/libs/metadata/src/reader/tables/field.rs

@@ -6,12 +6,12 @@ impl std::fmt::Debug for Field<'_> {
     }
 }
 
-impl Field<'_> {
+impl<'a> Field<'a> {
     pub fn flags(&self) -> FieldAttributes {
         FieldAttributes(self.usize(0).try_into().unwrap())
     }
 
-    pub fn name(&self) -> &str {
+    pub fn name(&self) -> &'a str {
         self.str(1)
     }
 

crates/libs/metadata/src/reader/tables/method_def.rs

@@ -6,7 +6,7 @@ impl std::fmt::Debug for MethodDef<'_> {
     }
 }
 
-impl MethodDef<'_> {
+impl<'a> MethodDef<'a> {
     pub fn rva(&self) -> usize {
         self.usize(0)
     }
@@ -19,7 +19,7 @@ impl MethodDef<'_> {
         MethodAttributes(self.usize(2).try_into().unwrap())
     }
 
-    pub fn name(&self) -> &str {
+    pub fn name(&self) -> &'a str {
         self.str(3)
     }
 

crates/libs/metadata/src/reader/tables/type_def.rs

@@ -6,16 +6,16 @@ impl std::fmt::Debug for TypeDef<'_> {
     }
 }
 
-impl TypeDef<'_> {
+impl<'a> TypeDef<'a> {
     pub fn flags(&self) -> TypeAttributes {
         TypeAttributes(self.usize(0).try_into().unwrap())
     }
 
-    pub fn name(&self) -> &str {
+    pub fn name(&self) -> &'a str {
         self.str(1)
     }
 
-    pub fn namespace(&self) -> &str {
+    pub fn namespace(&self) -> &'a str {
         self.str(2)
     }
 
@@ -27,11 +27,11 @@ impl TypeDef<'_> {
         Some(self.decode(3))
     }
 
-    pub fn fields(&self) -> RowIterator<Field> {
+    pub fn fields(&self) -> RowIterator<'a, Field<'a>> {
         self.list(4)
     }
 
-    pub fn methods(&self) -> RowIterator<MethodDef> {
+    pub fn methods(&self) -> RowIterator<'a, MethodDef<'a>> {
         self.list(5)
     }
 

crates/libs/metadata/src/reader/index.rs …s/libs/metadata/src/reader/type_index.rscrates/libs/metadata/src/reader/index.rs renamed to crates/libs/metadata/src/reader/type_index.rs crates/libs/metadata/src/reader/index.rs renamed to crates/libs/metadata/src/reader/type_index.rs

@@ -1,12 +1,12 @@
 use super::*;
 
-pub struct Index {
+pub struct TypeIndex {
     files: Vec<File>,
     types: HashMap<String, HashMap<String, Vec<(usize, usize)>>>,
     nested: HashMap<(usize, usize), Vec<usize>>,
 }
 
-impl Index {
+impl TypeIndex {
     pub fn read<P: AsRef<std::path::Path>>(path: P) -> Option<Self> {
         Some(Self::new(vec![File::read(path)?]))
     }
@@ -52,7 +52,21 @@ impl Index {
         &self.files[pos]
     }
 
-    pub fn all(&self) -> impl Iterator<Item = TypeDef> + '_ {
+    pub fn iter(&self) -> impl Iterator<Item = (&str, &str, TypeDef)> + '_ {
+        self.types
+            .iter()
+            .flat_map(|(namespace, types)| {
+                types
+                    .iter()
+                    .map(move |(name, types)| (namespace.as_str(), name.as_str(), types))
+            })
+            .flat_map(|(namespace, name, types)| types.iter().map(move |ty| (namespace, name, ty)))
+            .map(|(namespace, name, (file, pos))| {
+                (namespace, name, TypeDef(Row::new(self, *file, *pos)))
+            })
+    }
+
+    pub fn types(&self) -> impl Iterator<Item = TypeDef> + '_ {
         self.types
             .values()
             .flat_map(|types| types.values())
@@ -69,6 +83,7 @@ impl Index {
             .map(|(file, pos)| TypeDef(Row::new(self, *file, *pos)))
     }
 
+    #[track_caller]
     pub fn expect(&self, namespace: &str, name: &str) -> TypeDef {
         let mut iter = self.get(namespace, name);