span/
ast_id.rs

1//! `AstIdMap` allows to create stable IDs for "large" syntax nodes like items
2//! and macro calls.
3//!
4//! Specifically, it enumerates all items in a file and uses position of a an
5//! item as an ID. That way, id's don't change unless the set of items itself
6//! changes.
7//!
8//! These IDs are tricky. If one of them invalidates, its interned ID invalidates,
9//! and this can cause *a lot* to be recomputed. For example, if you invalidate the ID
10//! of a struct, and that struct has an impl (any impl!) this will cause the `Self`
11//! type of the impl to invalidate, which will cause the all impls queries to be
12//! invalidated, which will cause every trait solve query in this crate *and* all
13//! transitive reverse dependencies to be invalidated, which is pretty much the worst
14//! thing that can happen incrementality wise.
15//!
16//! So we want these IDs to stay as stable as possible. For top-level items, we store
17//! their kind and name, which should be unique, but since they can still not be, we
18//! also store an index disambiguator. For nested items, we also store the ID of their
19//! parent. For macro calls, we store the macro name and an index. There aren't usually
20//! a lot of macro calls in item position, and invalidation in bodies is not much of
21//! a problem, so this should be enough.
22
23use std::{
24    any::type_name,
25    fmt,
26    hash::{BuildHasher, Hash, Hasher},
27    marker::PhantomData,
28};
29
30use la_arena::{Arena, Idx, RawIdx};
31use rustc_hash::{FxBuildHasher, FxHashMap};
32use syntax::{
33    AstNode, AstPtr, SyntaxKind, SyntaxNode, SyntaxNodePtr,
34    ast::{self, HasName},
35    match_ast,
36};
37
38// The first index is always the root node's AstId
39/// The root ast id always points to the encompassing file, using this in spans is discouraged as
40/// any range relative to it will be effectively absolute, ruining the entire point of anchored
41/// relative text ranges.
42pub const ROOT_ERASED_FILE_AST_ID: ErasedFileAstId =
43    ErasedFileAstId(pack_hash_index_and_kind(0, 0, ErasedFileAstIdKind::Root as u32));
44
45/// ErasedFileAstId used as the span for syntax node fixups. Any Span containing this file id is to be
46/// considered fake.
47/// Do not modify this, it is used by the proc-macro server.
48pub const FIXUP_ERASED_FILE_AST_ID_MARKER: ErasedFileAstId =
49    ErasedFileAstId(pack_hash_index_and_kind(0, 0, ErasedFileAstIdKind::Fixup as u32));
50
51/// [`ErasedFileAstId`] used as the span for syntax nodes that should not be mapped down to
52/// macro expansion. Any `Span` containing this file id is to be considered fake.
53pub const NO_DOWNMAP_ERASED_FILE_AST_ID_MARKER: ErasedFileAstId =
54    ErasedFileAstId(pack_hash_index_and_kind(0, 0, ErasedFileAstIdKind::NoDownmap as u32));
55
56/// This is a type erased FileAstId.
57#[derive(Clone, Copy, PartialEq, Eq, Hash)]
58pub struct ErasedFileAstId(u32);
59
60impl fmt::Debug for ErasedFileAstId {
61    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
62        let kind = self.kind();
63        macro_rules! kind {
64            ($($kind:ident),* $(,)?) => {
65                if false {
66                    // Ensure we covered all variants.
67                    match ErasedFileAstIdKind::Root {
68                        $( ErasedFileAstIdKind::$kind => {} )*
69                    }
70                    unreachable!()
71                }
72                $( else if kind == ErasedFileAstIdKind::$kind as u32 {
73                    stringify!($kind)
74                } )*
75                else {
76                    "Unknown"
77                }
78            };
79        }
80        let kind = kind!(
81            Root,
82            Enum,
83            Struct,
84            Union,
85            ExternCrate,
86            MacroDef,
87            MacroRules,
88            Module,
89            Static,
90            Trait,
91            TraitAlias,
92            Variant,
93            Const,
94            Fn,
95            MacroCall,
96            TypeAlias,
97            ExternBlock,
98            Use,
99            Impl,
100            BlockExpr,
101            AsmExpr,
102            Fixup,
103            NoDownmap,
104        );
105        if f.alternate() {
106            write!(f, "{kind}[{:04X}, {}]", self.hash_value(), self.index())
107        } else {
108            f.debug_struct("ErasedFileAstId")
109                .field("kind", &format_args!("{kind}"))
110                .field("index", &self.index())
111                .field("hash", &format_args!("{:04X}", self.hash_value()))
112                .finish()
113        }
114    }
115}
116
117#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
118#[repr(u8)]
119enum ErasedFileAstIdKind {
120    /// This needs to not change because it's depended upon by the proc macro server.
121    Fixup = 0,
122    // The following are associated with `ErasedHasNameFileAstId`.
123    Enum,
124    Struct,
125    Union,
126    ExternCrate,
127    MacroDef,
128    MacroRules,
129    Module,
130    Static,
131    Trait,
132    TraitAlias,
133    // Until here associated with `ErasedHasNameFileAstId`.
134    // The following are associated with `ErasedAssocItemFileAstId`.
135    Variant,
136    Const,
137    Fn,
138    MacroCall,
139    TypeAlias,
140    // Until here associated with `ErasedAssocItemFileAstId`.
141    // Extern blocks don't really have any identifying property unfortunately.
142    ExternBlock,
143    // FIXME: If we store the final `UseTree` instead of the top-level `Use`, we can store its name,
144    // and be way more granular for incrementality, at the expense of increased memory usage.
145    // Use IDs aren't used a lot. The main thing that stores them is the def map. So everything that
146    // uses the def map will be invalidated. That includes infers, and so is pretty bad, but our
147    // def map incrementality story is pretty bad anyway and needs to be improved (see
148    // https://rust-lang.zulipchat.com/#narrow/channel/185405-t-compiler.2Frust-analyzer/topic/.60infer.60.20queries.20and.20splitting.20.60DefMap.60).
149    // So I left this as-is for now, as the def map improvement should also mitigate this.
150    Use,
151    /// Associated with [`ImplFileAstId`].
152    Impl,
153    /// Associated with [`BlockExprFileAstId`].
154    BlockExpr,
155    // `global_asm!()` is an item, so we need to give it an `AstId`. So we give to all inline asm
156    // because incrementality is not a problem, they will always be the only item in the macro file,
157    // and memory usage also not because they're rare.
158    AsmExpr,
159    /// Represents a fake [`ErasedFileAstId`] that should not be mapped down to macro expansion
160    /// result.
161    NoDownmap,
162    /// Keep this last.
163    Root,
164}
165
166// First hash, then index, then kind.
167const HASH_BITS: u32 = 16;
168const INDEX_BITS: u32 = 11;
169const KIND_BITS: u32 = 5;
170const _: () = assert!(ErasedFileAstIdKind::Root as u32 <= ((1 << KIND_BITS) - 1));
171const _: () = assert!(HASH_BITS + INDEX_BITS + KIND_BITS == u32::BITS);
172
173#[inline]
174const fn u16_hash(hash: u64) -> u16 {
175    // We do basically the same as `FxHasher`. We don't use rustc-hash and truncate because the
176    // higher bits have more entropy, but unlike rustc-hash we don't rotate because it rotates
177    // for hashmaps that just use the low bits, but we compare all bits.
178    const K: u16 = 0xecc5;
179    let (part1, part2, part3, part4) =
180        (hash as u16, (hash >> 16) as u16, (hash >> 32) as u16, (hash >> 48) as u16);
181    part1
182        .wrapping_add(part2)
183        .wrapping_mul(K)
184        .wrapping_add(part3)
185        .wrapping_mul(K)
186        .wrapping_add(part4)
187        .wrapping_mul(K)
188}
189
190#[inline]
191const fn pack_hash_index_and_kind(hash: u16, index: u32, kind: u32) -> u32 {
192    (hash as u32) | (index << HASH_BITS) | (kind << (HASH_BITS + INDEX_BITS))
193}
194
195impl ErasedFileAstId {
196    #[inline]
197    fn hash_value(self) -> u16 {
198        self.0 as u16
199    }
200
201    #[inline]
202    fn index(self) -> u32 {
203        (self.0 << KIND_BITS) >> (HASH_BITS + KIND_BITS)
204    }
205
206    #[inline]
207    fn kind(self) -> u32 {
208        self.0 >> (HASH_BITS + INDEX_BITS)
209    }
210
211    fn ast_id_for(
212        node: &SyntaxNode,
213        index_map: &mut ErasedAstIdNextIndexMap,
214        parent: Option<&ErasedFileAstId>,
215    ) -> Option<ErasedFileAstId> {
216        // Blocks are deliberately not here - we only want to allocate a block if it contains items.
217        has_name_ast_id(node, index_map)
218            .or_else(|| assoc_item_ast_id(node, index_map, parent))
219            .or_else(|| extern_block_ast_id(node, index_map))
220            .or_else(|| use_ast_id(node, index_map))
221            .or_else(|| impl_ast_id(node, index_map))
222            .or_else(|| asm_expr_ast_id(node, index_map))
223    }
224
225    fn should_alloc(node: &SyntaxNode) -> bool {
226        let kind = node.kind();
227        should_alloc_has_name(kind)
228            || should_alloc_assoc_item(kind)
229            || ast::ExternBlock::can_cast(kind)
230            || ast::Use::can_cast(kind)
231            || ast::Impl::can_cast(kind)
232            || ast::AsmExpr::can_cast(kind)
233    }
234
235    #[inline]
236    pub fn into_raw(self) -> u32 {
237        self.0
238    }
239
240    #[inline]
241    pub const fn from_raw(v: u32) -> Self {
242        Self(v)
243    }
244}
245
246pub trait AstIdNode: AstNode {}
247
248/// `AstId` points to an AST node in a specific file.
249pub struct FileAstId<N> {
250    raw: ErasedFileAstId,
251    _marker: PhantomData<fn() -> N>,
252}
253
254/// Traits are manually implemented because `derive` adds redundant bounds.
255impl<N> Clone for FileAstId<N> {
256    #[inline]
257    fn clone(&self) -> FileAstId<N> {
258        *self
259    }
260}
261impl<N> Copy for FileAstId<N> {}
262
263impl<N> PartialEq for FileAstId<N> {
264    fn eq(&self, other: &Self) -> bool {
265        self.raw == other.raw
266    }
267}
268impl<N> Eq for FileAstId<N> {}
269impl<N> Hash for FileAstId<N> {
270    fn hash<H: Hasher>(&self, hasher: &mut H) {
271        self.raw.hash(hasher);
272    }
273}
274
275impl<N> fmt::Debug for FileAstId<N> {
276    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
277        write!(f, "FileAstId::<{}>({:?})", type_name::<N>(), self.raw)
278    }
279}
280
281impl<N> FileAstId<N> {
282    // Can't make this a From implementation because of coherence
283    #[inline]
284    pub fn upcast<M: AstIdNode>(self) -> FileAstId<M>
285    where
286        N: Into<M>,
287    {
288        FileAstId { raw: self.raw, _marker: PhantomData }
289    }
290
291    #[inline]
292    pub fn erase(self) -> ErasedFileAstId {
293        self.raw
294    }
295}
296
297#[derive(Hash)]
298struct ErasedHasNameFileAstId<'a> {
299    name: &'a str,
300}
301
302/// This holds the ast ID for variants too (they're a kind of assoc item).
303#[derive(Hash)]
304struct ErasedAssocItemFileAstId<'a> {
305    /// Subtle: items in `extern` blocks **do not** store the ID of the extern block here.
306    /// Instead this is left empty. The reason is that `ExternBlockFileAstId` is pretty unstable
307    /// (it contains only an index), and extern blocks don't introduce a new scope, so storing
308    /// the extern block ID will do more harm to incrementality than help.
309    parent: Option<ErasedFileAstId>,
310    properties: ErasedHasNameFileAstId<'a>,
311}
312
313#[derive(Debug, Clone, PartialEq, Eq, Hash)]
314struct ImplFileAstId<'a> {
315    /// This can be `None` if the `Self` type is not a named type, or if it is inside a macro call.
316    self_ty_name: Option<&'a str>,
317    /// This can be `None` if this is an inherent impl, or if the trait name is inside a macro call.
318    trait_name: Option<&'a str>,
319}
320
321#[derive(Debug, Clone, PartialEq, Eq, Hash)]
322struct BlockExprFileAstId {
323    parent: Option<ErasedFileAstId>,
324}
325
326impl AstIdNode for ast::ExternBlock {}
327
328fn extern_block_ast_id(
329    node: &SyntaxNode,
330    index_map: &mut ErasedAstIdNextIndexMap,
331) -> Option<ErasedFileAstId> {
332    if ast::ExternBlock::can_cast(node.kind()) {
333        Some(index_map.new_id(ErasedFileAstIdKind::ExternBlock, ()))
334    } else {
335        None
336    }
337}
338
339impl AstIdNode for ast::Use {}
340
341fn use_ast_id(
342    node: &SyntaxNode,
343    index_map: &mut ErasedAstIdNextIndexMap,
344) -> Option<ErasedFileAstId> {
345    if ast::Use::can_cast(node.kind()) {
346        Some(index_map.new_id(ErasedFileAstIdKind::Use, ()))
347    } else {
348        None
349    }
350}
351
352impl AstIdNode for ast::AsmExpr {}
353
354fn asm_expr_ast_id(
355    node: &SyntaxNode,
356    index_map: &mut ErasedAstIdNextIndexMap,
357) -> Option<ErasedFileAstId> {
358    if ast::AsmExpr::can_cast(node.kind()) {
359        Some(index_map.new_id(ErasedFileAstIdKind::AsmExpr, ()))
360    } else {
361        None
362    }
363}
364
365impl AstIdNode for ast::Impl {}
366
367fn impl_ast_id(
368    node: &SyntaxNode,
369    index_map: &mut ErasedAstIdNextIndexMap,
370) -> Option<ErasedFileAstId> {
371    if let Some(node) = ast::Impl::cast(node.clone()) {
372        let type_as_name = |ty: Option<ast::Type>| match ty? {
373            ast::Type::PathType(it) => Some(it.path()?.segment()?.name_ref()?),
374            _ => None,
375        };
376        let self_ty_name = type_as_name(node.self_ty());
377        let trait_name = type_as_name(node.trait_());
378        let data = ImplFileAstId {
379            self_ty_name: self_ty_name.as_ref().map(|it| it.text_non_mutable()),
380            trait_name: trait_name.as_ref().map(|it| it.text_non_mutable()),
381        };
382        Some(index_map.new_id(ErasedFileAstIdKind::Impl, data))
383    } else {
384        None
385    }
386}
387
388// Blocks aren't `AstIdNode`s deliberately, because unlike other nodes, not all blocks get their own
389// ast id, only if they have items. To account for that we have a different, fallible, API for blocks.
390// impl !AstIdNode for ast::BlockExpr {}
391
392fn block_expr_ast_id(
393    node: &SyntaxNode,
394    index_map: &mut ErasedAstIdNextIndexMap,
395    parent: Option<&ErasedFileAstId>,
396) -> Option<ErasedFileAstId> {
397    if ast::BlockExpr::can_cast(node.kind()) {
398        Some(
399            index_map.new_id(
400                ErasedFileAstIdKind::BlockExpr,
401                BlockExprFileAstId { parent: parent.copied() },
402            ),
403        )
404    } else {
405        None
406    }
407}
408
409#[derive(Default)]
410struct ErasedAstIdNextIndexMap(FxHashMap<(ErasedFileAstIdKind, u16), u32>);
411
412impl ErasedAstIdNextIndexMap {
413    #[inline]
414    fn new_id(&mut self, kind: ErasedFileAstIdKind, data: impl Hash) -> ErasedFileAstId {
415        let hash = FxBuildHasher.hash_one(&data);
416        let initial_hash = u16_hash(hash);
417        // Even though 2^INDEX_BITS=2048 items with the same hash seems like a lot,
418        // it could happen with macro calls or `use`s in macro-generated files. So we want
419        // to handle it gracefully. We just increment the hash.
420        let mut hash = initial_hash;
421        let index = loop {
422            match self.0.entry((kind, hash)) {
423                std::collections::hash_map::Entry::Occupied(mut entry) => {
424                    let i = entry.get_mut();
425                    if *i < ((1 << INDEX_BITS) - 1) {
426                        *i += 1;
427                        break *i;
428                    }
429                }
430                std::collections::hash_map::Entry::Vacant(entry) => {
431                    entry.insert(0);
432                    break 0;
433                }
434            }
435            hash = hash.wrapping_add(1);
436            if hash == initial_hash {
437                // That's 2^27=134,217,728 items!
438                panic!("you have way too many items in the same file!");
439            }
440        };
441        let kind = kind as u32;
442        ErasedFileAstId(pack_hash_index_and_kind(hash, index, kind))
443    }
444}
445
446macro_rules! register_enum_ast_id {
447    (impl $AstIdNode:ident for $($ident:ident),+ ) => {
448        $(
449            impl $AstIdNode for ast::$ident {}
450        )+
451    };
452}
453register_enum_ast_id! {
454    impl AstIdNode for
455    Item, AnyHasGenericParams, Adt, Macro,
456    AssocItem
457}
458
459macro_rules! register_has_name_ast_id {
460    (impl $AstIdNode:ident for $($ident:ident = $name_method:ident),+ ) => {
461        $(
462            impl $AstIdNode for ast::$ident {}
463        )+
464
465        fn has_name_ast_id(node: &SyntaxNode, index_map: &mut ErasedAstIdNextIndexMap) -> Option<ErasedFileAstId> {
466            match_ast! {
467                match node {
468                    $(
469                        ast::$ident(node) => {
470                            let name = node.$name_method();
471                            let name = name.as_ref().map_or("", |it| it.text_non_mutable());
472                            let result = ErasedHasNameFileAstId {
473                                name,
474                            };
475                            Some(index_map.new_id(ErasedFileAstIdKind::$ident, result))
476                        },
477                    )*
478                    _ => None,
479                }
480            }
481        }
482
483        fn should_alloc_has_name(kind: SyntaxKind) -> bool {
484            false $( || ast::$ident::can_cast(kind) )*
485        }
486    };
487}
488register_has_name_ast_id! {
489    impl AstIdNode for
490        Enum = name,
491        Struct = name,
492        Union = name,
493        ExternCrate = name_ref,
494        MacroDef = name,
495        MacroRules = name,
496        Module = name,
497        Static = name,
498        Trait = name
499}
500
501macro_rules! register_assoc_item_ast_id {
502    (impl $AstIdNode:ident for $($ident:ident = $name_callback:expr),+ ) => {
503        $(
504            impl $AstIdNode for ast::$ident {}
505        )+
506
507        fn assoc_item_ast_id(
508            node: &SyntaxNode,
509            index_map: &mut ErasedAstIdNextIndexMap,
510            parent: Option<&ErasedFileAstId>,
511        ) -> Option<ErasedFileAstId> {
512            match_ast! {
513                match node {
514                    $(
515                        ast::$ident(node) => {
516                            let name = $name_callback(node);
517                            let name = name.as_ref().map_or("", |it| it.text_non_mutable());
518                            let properties = ErasedHasNameFileAstId {
519                                name,
520                            };
521                            let result = ErasedAssocItemFileAstId {
522                                parent: parent.copied(),
523                                properties,
524                            };
525                            Some(index_map.new_id(ErasedFileAstIdKind::$ident, result))
526                        },
527                    )*
528                    _ => None,
529                }
530            }
531        }
532
533        fn should_alloc_assoc_item(kind: SyntaxKind) -> bool {
534            false $( || ast::$ident::can_cast(kind) )*
535        }
536    };
537}
538register_assoc_item_ast_id! {
539    impl AstIdNode for
540    Variant = |it: ast::Variant| it.name(),
541    Const = |it: ast::Const| it.name(),
542    Fn = |it: ast::Fn| it.name(),
543    MacroCall = |it: ast::MacroCall| it.path().and_then(|path| path.segment()?.name_ref()),
544    TypeAlias = |it: ast::TypeAlias| it.name()
545}
546
547/// Maps items' `SyntaxNode`s to `ErasedFileAstId`s and back.
548#[derive(Default)]
549pub struct AstIdMap {
550    /// An arena of the ptrs and their associated ID.
551    arena: Arena<(SyntaxNodePtr, ErasedFileAstId)>,
552    /// Map ptr to id.
553    ptr_map: hashbrown::HashTable<ArenaId>,
554    /// Map id to ptr.
555    id_map: hashbrown::HashTable<ArenaId>,
556}
557
558type ArenaId = Idx<(SyntaxNodePtr, ErasedFileAstId)>;
559
560impl fmt::Debug for AstIdMap {
561    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
562        f.debug_struct("AstIdMap").field("arena", &self.arena).finish()
563    }
564}
565
566impl PartialEq for AstIdMap {
567    fn eq(&self, other: &Self) -> bool {
568        self.arena == other.arena
569    }
570}
571impl Eq for AstIdMap {}
572
573#[derive(Debug, Clone, Copy, PartialEq, Eq)]
574enum ContainsItems {
575    Yes,
576    No,
577}
578
579impl AstIdMap {
580    pub fn from_source(node: &SyntaxNode) -> AstIdMap {
581        assert!(node.parent().is_none());
582        let mut res = AstIdMap::default();
583        let mut index_map = ErasedAstIdNextIndexMap::default();
584
585        // Ensure we allocate the root.
586        res.arena.alloc((SyntaxNodePtr::new(node), ROOT_ERASED_FILE_AST_ID));
587
588        // By walking the tree in breadth-first order we make sure that parents
589        // get lower ids then children. That is, adding a new child does not
590        // change parent's id. This means that, say, adding a new function to a
591        // trait does not change ids of top-level items, which helps caching.
592
593        // This contains the stack of the `BlockExpr`s we are under. We do this
594        // so we only allocate `BlockExpr`s if they contain items.
595        // The general idea is: when we enter a block we push `(block, false)` here.
596        // Items inside the block are attributed to the block's container, not the block.
597        // For the first item we find inside a block, we make this `(block, true)`
598        // and create an ast id for the block. When exiting the block we pop it,
599        // whether or not we created an ast id for it.
600        // It may seem that with this setup we will generate an ID for blocks that
601        // have no items directly but have items inside other items inside them.
602        // This is true, but it doesn't matter, because such blocks can't exist.
603        // After all, the block will then contain the *outer* item, so we allocate
604        // an ID for it anyway.
605        let mut blocks = Vec::new();
606        let mut curr_layer = vec![(node.clone(), None)];
607        let mut next_layer = vec![];
608        while !curr_layer.is_empty() {
609            curr_layer.drain(..).for_each(|(node, parent_idx)| {
610                let mut preorder = node.preorder();
611                while let Some(event) = preorder.next() {
612                    match event {
613                        syntax::WalkEvent::Enter(node) => {
614                            if ast::BlockExpr::can_cast(node.kind()) {
615                                blocks.push((node, ContainsItems::No));
616                            } else if ErasedFileAstId::should_alloc(&node) {
617                                // Allocate blocks on-demand, only if they have items.
618                                // We don't associate items with blocks, only with items, since block IDs can be quite unstable.
619                                // FIXME: Is this the correct thing to do? Macro calls might actually be more incremental if
620                                // associated with blocks (not sure). Either way it's not a big deal.
621                                if let Some((
622                                    last_block_node,
623                                    already_allocated @ ContainsItems::No,
624                                )) = blocks.last_mut()
625                                {
626                                    let block_ast_id = block_expr_ast_id(
627                                        last_block_node,
628                                        &mut index_map,
629                                        parent_of(parent_idx, &res),
630                                    )
631                                    .expect("not a BlockExpr");
632                                    res.arena
633                                        .alloc((SyntaxNodePtr::new(last_block_node), block_ast_id));
634                                    *already_allocated = ContainsItems::Yes;
635                                }
636
637                                let parent = parent_of(parent_idx, &res);
638                                let ast_id =
639                                    ErasedFileAstId::ast_id_for(&node, &mut index_map, parent)
640                                        .expect("this node should have an ast id");
641                                let idx = res.arena.alloc((SyntaxNodePtr::new(&node), ast_id));
642
643                                next_layer.extend(node.children().map(|child| (child, Some(idx))));
644                                preorder.skip_subtree();
645                            }
646                        }
647                        syntax::WalkEvent::Leave(node) => {
648                            if ast::BlockExpr::can_cast(node.kind()) {
649                                assert_eq!(
650                                    blocks.pop().map(|it| it.0),
651                                    Some(node),
652                                    "left a BlockExpr we never entered"
653                                );
654                            }
655                        }
656                    }
657                }
658            });
659            std::mem::swap(&mut curr_layer, &mut next_layer);
660            assert!(blocks.is_empty(), "didn't leave all BlockExprs");
661        }
662
663        res.ptr_map = hashbrown::HashTable::with_capacity(res.arena.len());
664        res.id_map = hashbrown::HashTable::with_capacity(res.arena.len());
665        for (idx, (ptr, ast_id)) in res.arena.iter() {
666            let ptr_hash = hash_ptr(ptr);
667            let ast_id_hash = hash_ast_id(ast_id);
668            match res.ptr_map.entry(
669                ptr_hash,
670                |idx2| *idx2 == idx,
671                |&idx| hash_ptr(&res.arena[idx].0),
672            ) {
673                hashbrown::hash_table::Entry::Occupied(_) => unreachable!(),
674                hashbrown::hash_table::Entry::Vacant(entry) => {
675                    entry.insert(idx);
676                }
677            }
678            match res.id_map.entry(
679                ast_id_hash,
680                |idx2| *idx2 == idx,
681                |&idx| hash_ast_id(&res.arena[idx].1),
682            ) {
683                hashbrown::hash_table::Entry::Occupied(_) => unreachable!(),
684                hashbrown::hash_table::Entry::Vacant(entry) => {
685                    entry.insert(idx);
686                }
687            }
688        }
689        res.arena.shrink_to_fit();
690        return res;
691
692        fn parent_of(parent_idx: Option<ArenaId>, res: &AstIdMap) -> Option<&ErasedFileAstId> {
693            let mut parent = parent_idx.map(|parent_idx| &res.arena[parent_idx].1);
694            if parent.is_some_and(|parent| parent.kind() == ErasedFileAstIdKind::ExternBlock as u32)
695            {
696                // See the comment on `ErasedAssocItemFileAstId` for why is this.
697                // FIXME: Technically there could be an extern block inside another item, e.g.:
698                // ```
699                // fn foo() {
700                //     extern "C" {
701                //         fn bar();
702                //     }
703                // }
704                // ```
705                // Here we want to make `foo()` the parent of `bar()`, but we make it `None`.
706                // Shouldn't be a big deal though.
707                parent = None;
708            }
709            parent
710        }
711    }
712
713    /// The root node.
714    pub fn root(&self) -> SyntaxNodePtr {
715        self.arena[Idx::from_raw(RawIdx::from_u32(0))].0
716    }
717
718    pub fn ast_id<N: AstIdNode>(&self, item: &N) -> FileAstId<N> {
719        self.ast_id_for_ptr(AstPtr::new(item))
720    }
721
722    /// Blocks may not be allocated (if they have no items), so they have a different API.
723    pub fn ast_id_for_block(&self, block: &ast::BlockExpr) -> Option<FileAstId<ast::BlockExpr>> {
724        self.ast_id_for_ptr_for_block(AstPtr::new(block))
725    }
726
727    pub fn ast_id_for_ptr<N: AstIdNode>(&self, ptr: AstPtr<N>) -> FileAstId<N> {
728        let ptr = ptr.syntax_node_ptr();
729        FileAstId { raw: self.erased_ast_id(ptr), _marker: PhantomData }
730    }
731
732    /// Blocks may not be allocated (if they have no items), so they have a different API.
733    pub fn ast_id_for_ptr_for_block(
734        &self,
735        ptr: AstPtr<ast::BlockExpr>,
736    ) -> Option<FileAstId<ast::BlockExpr>> {
737        let ptr = ptr.syntax_node_ptr();
738        self.try_erased_ast_id(ptr).map(|raw| FileAstId { raw, _marker: PhantomData })
739    }
740
741    fn erased_ast_id(&self, ptr: SyntaxNodePtr) -> ErasedFileAstId {
742        self.try_erased_ast_id(ptr).unwrap_or_else(|| {
743            panic!(
744                "Can't find SyntaxNodePtr {:?} in AstIdMap:\n{:?}",
745                ptr,
746                self.arena.iter().map(|(_id, i)| i).collect::<Vec<_>>(),
747            )
748        })
749    }
750
751    fn try_erased_ast_id(&self, ptr: SyntaxNodePtr) -> Option<ErasedFileAstId> {
752        let hash = hash_ptr(&ptr);
753        let idx = *self.ptr_map.find(hash, |&idx| self.arena[idx].0 == ptr)?;
754        Some(self.arena[idx].1)
755    }
756
757    // Don't bound on `AstIdNode` here, because `BlockExpr`s are also valid here (`ast::BlockExpr`
758    // doesn't always have a matching `FileAstId`, but a `FileAstId<ast::BlockExpr>` always has
759    // a matching node).
760    pub fn get<N: AstNode>(&self, id: FileAstId<N>) -> AstPtr<N> {
761        let ptr = self.get_erased(id.raw);
762        AstPtr::try_from_raw(ptr)
763            .unwrap_or_else(|| panic!("AstIdMap node mismatch with node `{ptr:?}`"))
764    }
765
766    pub fn get_erased(&self, id: ErasedFileAstId) -> SyntaxNodePtr {
767        let hash = hash_ast_id(&id);
768        match self.id_map.find(hash, |&idx| self.arena[idx].1 == id) {
769            Some(&idx) => self.arena[idx].0,
770            None => panic!(
771                "Can't find ast id {:?} in AstIdMap:\n{:?}",
772                id,
773                self.arena.iter().map(|(_id, i)| i).collect::<Vec<_>>(),
774            ),
775        }
776    }
777}
778
779#[inline]
780fn hash_ptr(ptr: &SyntaxNodePtr) -> u64 {
781    FxBuildHasher.hash_one(ptr)
782}
783
784#[inline]
785fn hash_ast_id(ptr: &ErasedFileAstId) -> u64 {
786    FxBuildHasher.hash_one(ptr)
787}
788
789#[cfg(test)]
790mod tests {
791    use syntax::{AstNode, Edition, SourceFile, SyntaxKind, SyntaxNodePtr, WalkEvent, ast};
792
793    use crate::AstIdMap;
794
795    #[test]
796    fn check_all_nodes() {
797        let syntax = SourceFile::parse(
798            r#"
799extern crate foo;
800fn foo() {
801    union U {}
802}
803struct S;
804macro_rules! m {}
805macro m2() {}
806trait Trait {}
807impl Trait for S {}
808impl S {}
809impl m!() {}
810impl m2!() for m!() {}
811type T = i32;
812enum E {
813    V1(),
814    V2 {},
815    V3,
816}
817struct S; // duplicate
818extern "C" {
819    static S: i32;
820}
821static mut S: i32 = 0;
822const FOO: i32 = 0;
823        "#,
824            Edition::CURRENT,
825        )
826        .syntax_node();
827        let ast_id_map = AstIdMap::from_source(&syntax);
828        for node in syntax.preorder() {
829            let WalkEvent::Enter(node) = node else { continue };
830            if !matches!(
831                node.kind(),
832                SyntaxKind::EXTERN_CRATE
833                    | SyntaxKind::FN
834                    | SyntaxKind::UNION
835                    | SyntaxKind::STRUCT
836                    | SyntaxKind::MACRO_RULES
837                    | SyntaxKind::MACRO_DEF
838                    | SyntaxKind::MACRO_CALL
839                    | SyntaxKind::TRAIT
840                    | SyntaxKind::IMPL
841                    | SyntaxKind::TYPE_ALIAS
842                    | SyntaxKind::ENUM
843                    | SyntaxKind::VARIANT
844                    | SyntaxKind::EXTERN_BLOCK
845                    | SyntaxKind::STATIC
846                    | SyntaxKind::CONST
847            ) {
848                continue;
849            }
850            let ptr = SyntaxNodePtr::new(&node);
851            let ast_id = ast_id_map.erased_ast_id(ptr);
852            let turn_back = ast_id_map.get_erased(ast_id);
853            assert_eq!(ptr, turn_back);
854        }
855    }
856
857    #[test]
858    fn different_names_get_different_hashes() {
859        let syntax = SourceFile::parse(
860            r#"
861fn foo() {}
862fn bar() {}
863        "#,
864            Edition::CURRENT,
865        )
866        .syntax_node();
867        let ast_id_map = AstIdMap::from_source(&syntax);
868        let fns = syntax.descendants().filter_map(ast::Fn::cast).collect::<Vec<_>>();
869        let [foo_fn, bar_fn] = fns.as_slice() else {
870            panic!("not exactly 2 functions");
871        };
872        let foo_fn_id = ast_id_map.ast_id(foo_fn);
873        let bar_fn_id = ast_id_map.ast_id(bar_fn);
874        assert_ne!(foo_fn_id.raw.hash_value(), bar_fn_id.raw.hash_value(), "hashes are equal");
875    }
876
877    #[test]
878    fn different_parents_get_different_hashes() {
879        let syntax = SourceFile::parse(
880            r#"
881fn foo() {
882    m!();
883}
884fn bar() {
885    m!();
886}
887        "#,
888            Edition::CURRENT,
889        )
890        .syntax_node();
891        let ast_id_map = AstIdMap::from_source(&syntax);
892        let macro_calls = syntax.descendants().filter_map(ast::MacroCall::cast).collect::<Vec<_>>();
893        let [macro_call_foo, macro_call_bar] = macro_calls.as_slice() else {
894            panic!("not exactly 2 macro calls");
895        };
896        let macro_call_foo_id = ast_id_map.ast_id(macro_call_foo);
897        let macro_call_bar_id = ast_id_map.ast_id(macro_call_bar);
898        assert_ne!(
899            macro_call_foo_id.raw.hash_value(),
900            macro_call_bar_id.raw.hash_value(),
901            "hashes are equal"
902        );
903    }
904
905    #[test]
906    fn blocks_with_no_items_have_no_id() {
907        let syntax = SourceFile::parse(
908            r#"
909fn foo() {
910    let foo = 1;
911    bar(foo);
912}
913        "#,
914            Edition::CURRENT,
915        )
916        .syntax_node();
917        let ast_id_map = AstIdMap::from_source(&syntax);
918        let block = syntax.descendants().find_map(ast::BlockExpr::cast).expect("no block");
919        assert!(ast_id_map.ast_id_for_block(&block).is_none());
920    }
921}