rustpython_ruff_python_ast/
name.rs

1use std::borrow::{Borrow, Cow};
2use std::fmt::{Debug, Display, Formatter, Write};
3use std::hash::{Hash, Hasher};
4use std::ops::Deref;
5
6use crate::Expr;
7use crate::generated::ExprName;
8
9#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
10#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
11#[cfg_attr(feature = "cache", derive(ruff_macros::CacheKey))]
12#[cfg_attr(feature = "salsa", derive(salsa::Update))]
13#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
14#[cfg_attr(
15    feature = "schemars",
16    derive(schemars::JsonSchema),
17    schemars(with = "String")
18)]
19pub struct Name(compact_str::CompactString);
20
21impl Name {
22    #[inline]
23    pub fn empty() -> Self {
24        Self(compact_str::CompactString::default())
25    }
26
27    #[inline]
28    pub fn new(name: impl AsRef<str>) -> Self {
29        Self(compact_str::CompactString::new(name))
30    }
31
32    #[inline]
33    pub const fn new_static(name: &'static str) -> Self {
34        Self(compact_str::CompactString::const_new(name))
35    }
36
37    pub fn shrink_to_fit(&mut self) {
38        self.0.shrink_to_fit();
39    }
40
41    pub fn as_str(&self) -> &str {
42        self.0.as_str()
43    }
44
45    pub fn push_str(&mut self, s: &str) {
46        self.0.push_str(s);
47    }
48}
49
50impl Debug for Name {
51    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
52        write!(f, "Name({:?})", self.as_str())
53    }
54}
55
56impl std::fmt::Write for Name {
57    fn write_str(&mut self, s: &str) -> std::fmt::Result {
58        self.0.push_str(s);
59        Ok(())
60    }
61}
62
63impl AsRef<str> for Name {
64    #[inline]
65    fn as_ref(&self) -> &str {
66        self.as_str()
67    }
68}
69
70impl Deref for Name {
71    type Target = str;
72
73    #[inline]
74    fn deref(&self) -> &Self::Target {
75        self.as_str()
76    }
77}
78
79impl Borrow<str> for Name {
80    #[inline]
81    fn borrow(&self) -> &str {
82        self.as_str()
83    }
84}
85
86impl<'a> From<&'a str> for Name {
87    #[inline]
88    fn from(s: &'a str) -> Self {
89        Name(s.into())
90    }
91}
92
93impl From<String> for Name {
94    #[inline]
95    fn from(s: String) -> Self {
96        Name(s.into())
97    }
98}
99
100impl<'a> From<&'a String> for Name {
101    #[inline]
102    fn from(s: &'a String) -> Self {
103        Name(s.into())
104    }
105}
106
107impl<'a> From<Cow<'a, str>> for Name {
108    #[inline]
109    fn from(cow: Cow<'a, str>) -> Self {
110        Name(cow.into())
111    }
112}
113
114impl From<Box<str>> for Name {
115    #[inline]
116    fn from(b: Box<str>) -> Self {
117        Name(b.into())
118    }
119}
120
121impl From<compact_str::CompactString> for Name {
122    #[inline]
123    fn from(value: compact_str::CompactString) -> Self {
124        Self(value)
125    }
126}
127
128impl From<Name> for compact_str::CompactString {
129    #[inline]
130    fn from(name: Name) -> Self {
131        name.0
132    }
133}
134
135impl From<Name> for String {
136    #[inline]
137    fn from(name: Name) -> Self {
138        name.as_str().into()
139    }
140}
141
142impl FromIterator<char> for Name {
143    fn from_iter<I: IntoIterator<Item = char>>(iter: I) -> Self {
144        Self(iter.into_iter().collect())
145    }
146}
147
148impl std::fmt::Display for Name {
149    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
150        f.write_str(self.as_str())
151    }
152}
153
154impl PartialEq<str> for Name {
155    #[inline]
156    fn eq(&self, other: &str) -> bool {
157        self.as_str() == other
158    }
159}
160
161impl PartialEq<Name> for str {
162    #[inline]
163    fn eq(&self, other: &Name) -> bool {
164        other == self
165    }
166}
167
168impl PartialEq<&str> for Name {
169    #[inline]
170    fn eq(&self, other: &&str) -> bool {
171        self.as_str() == *other
172    }
173}
174
175impl PartialEq<Name> for &str {
176    #[inline]
177    fn eq(&self, other: &Name) -> bool {
178        other == self
179    }
180}
181
182impl PartialEq<String> for Name {
183    fn eq(&self, other: &String) -> bool {
184        self == other.as_str()
185    }
186}
187
188impl PartialEq<Name> for String {
189    #[inline]
190    fn eq(&self, other: &Name) -> bool {
191        other == self
192    }
193}
194
195impl PartialEq<&String> for Name {
196    #[inline]
197    fn eq(&self, other: &&String) -> bool {
198        self.as_str() == *other
199    }
200}
201
202impl PartialEq<Name> for &String {
203    #[inline]
204    fn eq(&self, other: &Name) -> bool {
205        other == self
206    }
207}
208
209/// A representation of a qualified name, like `typing.List`.
210#[derive(Debug, Clone, PartialEq, Eq, Hash)]
211pub struct QualifiedName<'a>(SegmentsVec<'a>);
212
213impl<'a> QualifiedName<'a> {
214    /// Create a [`QualifiedName`] from a dotted name.
215    ///
216    /// ```rust
217    /// # use ruff_python_ast::name::QualifiedName;
218    ///
219    /// assert_eq!(QualifiedName::from_dotted_name("typing.List").segments(), ["typing", "List"]);
220    /// assert_eq!(QualifiedName::from_dotted_name("list").segments(), ["", "list"]);
221    /// ```
222    #[inline]
223    pub fn from_dotted_name(name: &'a str) -> Self {
224        if let Some(dot) = name.find('.') {
225            let mut builder = QualifiedNameBuilder::default();
226            builder.push(&name[..dot]);
227            builder.extend(name[dot + 1..].split('.'));
228            builder.build()
229        } else {
230            Self::builtin(name)
231        }
232    }
233
234    /// Creates a name that's guaranteed not be a built in
235    #[inline]
236    pub fn user_defined(name: &'a str) -> Self {
237        name.split('.').collect()
238    }
239
240    /// Creates a qualified name for a built in
241    #[inline]
242    pub fn builtin(name: &'a str) -> Self {
243        debug_assert!(!name.contains('.'));
244        Self(SegmentsVec::Stack(SegmentsStack::from_slice(&["", name])))
245    }
246
247    #[inline]
248    pub fn segments(&self) -> &[&'a str] {
249        self.0.as_slice()
250    }
251
252    /// If the first segment is empty, the `CallPath` represents a "builtin binding".
253    ///
254    /// A builtin binding is the binding that a symbol has if it was part of Python's
255    /// global scope without any imports taking place. However, if builtin members are
256    /// accessed explicitly via the `builtins` module, they will not have a
257    /// "builtin binding", so this method will return `false`.
258    ///
259    /// Ex) `["", "bool"]` -> `"bool"`
260    fn is_builtin(&self) -> bool {
261        matches!(self.segments(), ["", ..])
262    }
263
264    /// If the call path is dot-prefixed, it's an unresolved relative import.
265    /// Ex) `[".foo", "bar"]` -> `".foo.bar"`
266    pub fn is_unresolved_import(&self) -> bool {
267        matches!(self.segments(), [".", ..])
268    }
269
270    pub fn starts_with(&self, other: &QualifiedName<'_>) -> bool {
271        self.segments().starts_with(other.segments())
272    }
273
274    /// Appends a member to the qualified name.
275    #[must_use]
276    pub fn append_member(self, member: &'a str) -> Self {
277        let mut inner = self.0;
278        inner.push(member);
279        Self(inner)
280    }
281
282    /// Extends the qualified name using the given members.
283    #[must_use]
284    pub fn extend_members<T: IntoIterator<Item = &'a str>>(self, members: T) -> Self {
285        let mut inner = self.0;
286        inner.extend(members);
287        Self(inner)
288    }
289}
290
291impl Display for QualifiedName<'_> {
292    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
293        let segments = self.segments();
294
295        if self.is_unresolved_import() {
296            let mut iter = segments.iter();
297            for segment in iter.by_ref() {
298                if *segment == "." {
299                    f.write_char('.')?;
300                } else {
301                    f.write_str(segment)?;
302                    break;
303                }
304            }
305            for segment in iter {
306                f.write_char('.')?;
307                f.write_str(segment)?;
308            }
309        } else {
310            let segments = if self.is_builtin() {
311                &segments[1..]
312            } else {
313                segments
314            };
315
316            let mut first = true;
317            for segment in segments {
318                if !first {
319                    f.write_char('.')?;
320                }
321
322                f.write_str(segment)?;
323                first = false;
324            }
325        }
326
327        Ok(())
328    }
329}
330
331impl<'a> FromIterator<&'a str> for QualifiedName<'a> {
332    fn from_iter<T: IntoIterator<Item = &'a str>>(iter: T) -> Self {
333        Self(SegmentsVec::from_iter(iter))
334    }
335}
336
337#[derive(Debug, Clone, Default)]
338pub struct QualifiedNameBuilder<'a> {
339    segments: SegmentsVec<'a>,
340}
341
342impl<'a> QualifiedNameBuilder<'a> {
343    pub fn with_capacity(capacity: usize) -> Self {
344        Self {
345            segments: SegmentsVec::with_capacity(capacity),
346        }
347    }
348
349    #[inline]
350    pub fn is_empty(&self) -> bool {
351        self.segments.is_empty()
352    }
353
354    #[inline]
355    pub fn push(&mut self, segment: &'a str) {
356        self.segments.push(segment);
357    }
358
359    #[inline]
360    pub fn pop(&mut self) {
361        self.segments.pop();
362    }
363
364    #[inline]
365    pub fn extend(&mut self, segments: impl IntoIterator<Item = &'a str>) {
366        self.segments.extend(segments);
367    }
368
369    #[inline]
370    pub fn extend_from_slice(&mut self, segments: &[&'a str]) {
371        self.segments.extend_from_slice(segments);
372    }
373
374    pub fn build(self) -> QualifiedName<'a> {
375        QualifiedName(self.segments)
376    }
377}
378
379#[derive(Debug, Clone, PartialEq, Eq, Hash)]
380pub struct UnqualifiedName<'a>(SegmentsVec<'a>);
381
382impl<'a> UnqualifiedName<'a> {
383    /// Convert an `Expr` to its [`UnqualifiedName`] (like `["typing", "List"]`).
384    pub fn from_expr(expr: &'a Expr) -> Option<Self> {
385        // Unroll the loop up to eight times, to match the maximum number of expected attributes.
386        // In practice, unrolling appears to give about a 4x speed-up on this hot path.
387        let attr1 = match expr {
388            Expr::Attribute(attr1) => attr1,
389            // Ex) `foo`
390            Expr::Name(ExprName { id, .. }) => return Some(Self::from_slice(&[id.as_str()])),
391            _ => return None,
392        };
393
394        let attr2 = match attr1.value.as_ref() {
395            Expr::Attribute(attr2) => attr2,
396            // Ex) `foo.bar`
397            Expr::Name(ExprName { id, .. }) => {
398                return Some(Self::from_slice(&[id.as_str(), attr1.attr.as_str()]));
399            }
400            _ => return None,
401        };
402
403        let attr3 = match attr2.value.as_ref() {
404            Expr::Attribute(attr3) => attr3,
405            // Ex) `foo.bar.baz`
406            Expr::Name(ExprName { id, .. }) => {
407                return Some(Self::from_slice(&[
408                    id.as_str(),
409                    attr2.attr.as_str(),
410                    attr1.attr.as_str(),
411                ]));
412            }
413            _ => return None,
414        };
415
416        let attr4 = match attr3.value.as_ref() {
417            Expr::Attribute(attr4) => attr4,
418            // Ex) `foo.bar.baz.bop`
419            Expr::Name(ExprName { id, .. }) => {
420                return Some(Self::from_slice(&[
421                    id.as_str(),
422                    attr3.attr.as_str(),
423                    attr2.attr.as_str(),
424                    attr1.attr.as_str(),
425                ]));
426            }
427            _ => return None,
428        };
429
430        let attr5 = match attr4.value.as_ref() {
431            Expr::Attribute(attr5) => attr5,
432            // Ex) `foo.bar.baz.bop.bap`
433            Expr::Name(ExprName { id, .. }) => {
434                return Some(Self::from_slice(&[
435                    id.as_str(),
436                    attr4.attr.as_str(),
437                    attr3.attr.as_str(),
438                    attr2.attr.as_str(),
439                    attr1.attr.as_str(),
440                ]));
441            }
442            _ => return None,
443        };
444
445        let attr6 = match attr5.value.as_ref() {
446            Expr::Attribute(attr6) => attr6,
447            // Ex) `foo.bar.baz.bop.bap.bab`
448            Expr::Name(ExprName { id, .. }) => {
449                return Some(Self::from_slice(&[
450                    id.as_str(),
451                    attr5.attr.as_str(),
452                    attr4.attr.as_str(),
453                    attr3.attr.as_str(),
454                    attr2.attr.as_str(),
455                    attr1.attr.as_str(),
456                ]));
457            }
458            _ => return None,
459        };
460
461        let attr7 = match attr6.value.as_ref() {
462            Expr::Attribute(attr7) => attr7,
463            // Ex) `foo.bar.baz.bop.bap.bab.bob`
464            Expr::Name(ExprName { id, .. }) => {
465                return Some(Self::from_slice(&[
466                    id.as_str(),
467                    attr6.attr.as_str(),
468                    attr5.attr.as_str(),
469                    attr4.attr.as_str(),
470                    attr3.attr.as_str(),
471                    attr2.attr.as_str(),
472                    attr1.attr.as_str(),
473                ]));
474            }
475            _ => return None,
476        };
477
478        let attr8 = match attr7.value.as_ref() {
479            Expr::Attribute(attr8) => attr8,
480            // Ex) `foo.bar.baz.bop.bap.bab.bob.bib`
481            Expr::Name(ExprName { id, .. }) => {
482                return Some(Self(SegmentsVec::from([
483                    id.as_str(),
484                    attr7.attr.as_str(),
485                    attr6.attr.as_str(),
486                    attr5.attr.as_str(),
487                    attr4.attr.as_str(),
488                    attr3.attr.as_str(),
489                    attr2.attr.as_str(),
490                    attr1.attr.as_str(),
491                ])));
492            }
493            _ => return None,
494        };
495
496        let mut segments = Vec::with_capacity(SMALL_LEN * 2);
497
498        let mut current = &*attr8.value;
499
500        loop {
501            current = match current {
502                Expr::Attribute(attr) => {
503                    segments.push(attr.attr.as_str());
504                    &*attr.value
505                }
506                Expr::Name(ExprName { id, .. }) => {
507                    segments.push(id.as_str());
508                    break;
509                }
510                _ => {
511                    return None;
512                }
513            }
514        }
515
516        segments.reverse();
517
518        // Append the attributes we visited before calling into the recursion.
519        segments.extend_from_slice(&[
520            attr8.attr.as_str(),
521            attr7.attr.as_str(),
522            attr6.attr.as_str(),
523            attr5.attr.as_str(),
524            attr4.attr.as_str(),
525            attr3.attr.as_str(),
526            attr2.attr.as_str(),
527            attr1.attr.as_str(),
528        ]);
529
530        Some(Self(SegmentsVec::from(segments)))
531    }
532
533    #[inline]
534    pub fn from_slice(segments: &[&'a str]) -> Self {
535        Self(SegmentsVec::from_slice(segments))
536    }
537
538    pub fn segments(&self) -> &[&'a str] {
539        self.0.as_slice()
540    }
541}
542
543impl Display for UnqualifiedName<'_> {
544    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
545        let mut first = true;
546        for segment in self.segments() {
547            if !first {
548                f.write_char('.')?;
549            }
550
551            f.write_str(segment)?;
552            first = false;
553        }
554
555        Ok(())
556    }
557}
558
559impl<'a> FromIterator<&'a str> for UnqualifiedName<'a> {
560    #[inline]
561    fn from_iter<T: IntoIterator<Item = &'a str>>(iter: T) -> Self {
562        Self(iter.into_iter().collect())
563    }
564}
565
566/// A smallvec like storage for qualified and unqualified name segments.
567///
568/// Stores up to 8 segments inline, and falls back to a heap-allocated vector for names with more segments.
569///
570/// ## Note
571/// The implementation doesn't use `SmallVec` v1 because its type definition has a variance problem.
572/// The incorrect variance leads the lifetime inference in the `SemanticModel` astray, causing
573/// all sort of "strange" lifetime errors. We can switch back to `SmallVec` when v2 is released.
574#[derive(Clone)]
575enum SegmentsVec<'a> {
576    Stack(SegmentsStack<'a>),
577    Heap(Vec<&'a str>),
578}
579
580impl<'a> SegmentsVec<'a> {
581    /// Creates an empty segment vec.
582    fn new() -> Self {
583        Self::Stack(SegmentsStack::default())
584    }
585
586    /// Creates a segment vec that has reserved storage for up to `capacity` items.
587    fn with_capacity(capacity: usize) -> Self {
588        if capacity <= SMALL_LEN {
589            Self::new()
590        } else {
591            Self::Heap(Vec::with_capacity(capacity))
592        }
593    }
594
595    #[cfg(test)]
596    const fn is_spilled(&self) -> bool {
597        matches!(self, SegmentsVec::Heap(_))
598    }
599
600    /// Initializes the segments from a slice.
601    #[inline]
602    fn from_slice(slice: &[&'a str]) -> Self {
603        if slice.len() <= SMALL_LEN {
604            SegmentsVec::Stack(SegmentsStack::from_slice(slice))
605        } else {
606            SegmentsVec::Heap(slice.to_vec())
607        }
608    }
609
610    /// Returns the segments as a slice.
611    #[inline]
612    fn as_slice(&self) -> &[&'a str] {
613        match self {
614            Self::Stack(stack) => stack.as_slice(),
615            Self::Heap(heap) => heap.as_slice(),
616        }
617    }
618
619    /// Pushes `name` to the end of the segments.
620    ///
621    /// Spills to the heap if the segments are stored on the stack and the 9th segment is pushed.
622    #[inline]
623    fn push(&mut self, name: &'a str) {
624        match self {
625            SegmentsVec::Stack(stack) => {
626                if let Err(segments) = stack.push(name) {
627                    *self = SegmentsVec::Heap(segments);
628                }
629            }
630            SegmentsVec::Heap(heap) => {
631                heap.push(name);
632            }
633        }
634    }
635
636    /// Pops the last segment from the end and returns it.
637    ///
638    /// Returns `None` if the vector is empty.
639    #[inline]
640    fn pop(&mut self) -> Option<&'a str> {
641        match self {
642            SegmentsVec::Stack(stack) => stack.pop(),
643            SegmentsVec::Heap(heap) => heap.pop(),
644        }
645    }
646
647    #[inline]
648    fn extend_from_slice(&mut self, slice: &[&'a str]) {
649        match self {
650            SegmentsVec::Stack(stack) => {
651                if let Err(segments) = stack.extend_from_slice(slice) {
652                    *self = SegmentsVec::Heap(segments);
653                }
654            }
655            SegmentsVec::Heap(heap) => heap.extend_from_slice(slice),
656        }
657    }
658}
659
660impl Default for SegmentsVec<'_> {
661    fn default() -> Self {
662        Self::new()
663    }
664}
665
666impl Debug for SegmentsVec<'_> {
667    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
668        f.debug_list().entries(self.as_slice()).finish()
669    }
670}
671
672impl<'a> Deref for SegmentsVec<'a> {
673    type Target = [&'a str];
674    fn deref(&self) -> &Self::Target {
675        self.as_slice()
676    }
677}
678
679impl<'b> PartialEq<SegmentsVec<'b>> for SegmentsVec<'_> {
680    fn eq(&self, other: &SegmentsVec<'b>) -> bool {
681        self.as_slice() == other.as_slice()
682    }
683}
684
685impl Eq for SegmentsVec<'_> {}
686
687impl Hash for SegmentsVec<'_> {
688    fn hash<H: Hasher>(&self, state: &mut H) {
689        self.as_slice().hash(state);
690    }
691}
692
693impl<'a> FromIterator<&'a str> for SegmentsVec<'a> {
694    #[inline]
695    fn from_iter<T: IntoIterator<Item = &'a str>>(iter: T) -> Self {
696        let mut segments = SegmentsVec::default();
697        segments.extend(iter);
698        segments
699    }
700}
701
702impl<'a> From<[&'a str; 8]> for SegmentsVec<'a> {
703    #[inline]
704    fn from(segments: [&'a str; 8]) -> Self {
705        SegmentsVec::Stack(SegmentsStack {
706            segments,
707            len: segments.len(),
708        })
709    }
710}
711
712impl<'a> From<Vec<&'a str>> for SegmentsVec<'a> {
713    #[inline]
714    fn from(segments: Vec<&'a str>) -> Self {
715        SegmentsVec::Heap(segments)
716    }
717}
718
719impl<'a> Extend<&'a str> for SegmentsVec<'a> {
720    #[inline]
721    fn extend<T: IntoIterator<Item = &'a str>>(&mut self, iter: T) {
722        match self {
723            SegmentsVec::Stack(stack) => {
724                if let Err(segments) = stack.extend(iter) {
725                    *self = SegmentsVec::Heap(segments);
726                }
727            }
728            SegmentsVec::Heap(heap) => {
729                heap.extend(iter);
730            }
731        }
732    }
733}
734
735const SMALL_LEN: usize = 8;
736
737#[derive(Debug, Clone, Default)]
738struct SegmentsStack<'a> {
739    segments: [&'a str; SMALL_LEN],
740    len: usize,
741}
742
743impl<'a> SegmentsStack<'a> {
744    #[inline]
745    fn from_slice(slice: &[&'a str]) -> Self {
746        assert!(slice.len() <= SMALL_LEN);
747
748        let mut segments: [&'a str; SMALL_LEN] = Default::default();
749        segments[..slice.len()].copy_from_slice(slice);
750        SegmentsStack {
751            segments,
752            len: slice.len(),
753        }
754    }
755
756    const fn capacity(&self) -> usize {
757        SMALL_LEN - self.len
758    }
759
760    #[inline]
761    fn as_slice(&self) -> &[&'a str] {
762        &self.segments[..self.len]
763    }
764
765    /// Pushes `name` to the end of the segments.
766    ///
767    /// Returns `Err` with a `Vec` containing all items (including `name`) if there's not enough capacity to push the name.
768    #[inline]
769    fn push(&mut self, name: &'a str) -> Result<(), Vec<&'a str>> {
770        if self.len < self.segments.len() {
771            self.segments[self.len] = name;
772            self.len += 1;
773            Ok(())
774        } else {
775            let mut segments = Vec::with_capacity(self.len * 2);
776            segments.extend_from_slice(&self.segments);
777            segments.push(name);
778            Err(segments)
779        }
780    }
781
782    /// Reserves spaces for `additional` segments.
783    ///
784    /// Returns `Err` with a `Vec` containing all segments and a capacity to store `additional` segments if
785    /// the stack needs to spill over to the heap to store `additional` segments.
786    #[inline]
787    fn reserve(&mut self, additional: usize) -> Result<(), Vec<&'a str>> {
788        if self.capacity() >= additional {
789            Ok(())
790        } else {
791            let mut segments = Vec::with_capacity(self.len + additional);
792            segments.extend_from_slice(self.as_slice());
793            Err(segments)
794        }
795    }
796
797    #[inline]
798    fn pop(&mut self) -> Option<&'a str> {
799        if self.len > 0 {
800            self.len -= 1;
801            Some(self.segments[self.len])
802        } else {
803            None
804        }
805    }
806
807    /// Extends the segments by appending `slice` to the end.
808    ///
809    /// Returns `Err` with a `Vec` containing all segments and the segments in `slice` if there's not enough capacity to append the names.
810    #[inline]
811    fn extend_from_slice(&mut self, slice: &[&'a str]) -> Result<(), Vec<&'a str>> {
812        let new_len = self.len + slice.len();
813        if slice.len() <= self.capacity() {
814            self.segments[self.len..new_len].copy_from_slice(slice);
815            self.len = new_len;
816            Ok(())
817        } else {
818            let mut segments = Vec::with_capacity(new_len);
819            segments.extend_from_slice(self.as_slice());
820            segments.extend_from_slice(slice);
821            Err(segments)
822        }
823    }
824
825    #[inline]
826    fn extend<I>(&mut self, iter: I) -> Result<(), Vec<&'a str>>
827    where
828        I: IntoIterator<Item = &'a str>,
829    {
830        let mut iter = iter.into_iter();
831        let (lower, _) = iter.size_hint();
832
833        // There's not enough space to store the lower bound of items. Spill to the heap.
834        if let Err(mut segments) = self.reserve(lower) {
835            segments.extend(iter);
836            return Err(segments);
837        }
838
839        // Copy over up to capacity items
840        for name in iter.by_ref().take(self.capacity()) {
841            self.segments[self.len] = name;
842            self.len += 1;
843        }
844
845        let Some(item) = iter.next() else {
846            // There are no more items to copy over and they all fit into capacity.
847            return Ok(());
848        };
849
850        // There are more items and there's not enough capacity to store them on the stack.
851        // Spill over to the heap and append the remaining items.
852        let mut segments = Vec::with_capacity(self.len * 2);
853        segments.extend_from_slice(self.as_slice());
854        segments.push(item);
855        segments.extend(iter);
856        Err(segments)
857    }
858}
859
860#[cfg(test)]
861mod tests {
862    use crate::name::SegmentsVec;
863
864    #[test]
865    fn empty_vec() {
866        let empty = SegmentsVec::new();
867        assert_eq!(empty.as_slice(), &[] as &[&str]);
868        assert!(!empty.is_spilled());
869    }
870
871    #[test]
872    fn from_slice_stack() {
873        let stack = SegmentsVec::from_slice(&["a", "b", "c"]);
874
875        assert_eq!(stack.as_slice(), &["a", "b", "c"]);
876        assert!(!stack.is_spilled());
877    }
878
879    #[test]
880    fn from_slice_stack_capacity() {
881        let stack = SegmentsVec::from_slice(&["a", "b", "c", "d", "e", "f", "g", "h"]);
882
883        assert_eq!(stack.as_slice(), &["a", "b", "c", "d", "e", "f", "g", "h"]);
884        assert!(!stack.is_spilled());
885    }
886
887    #[test]
888    fn from_slice_heap() {
889        let heap = SegmentsVec::from_slice(&["a", "b", "c", "d", "e", "f", "g", "h", "i"]);
890
891        assert_eq!(
892            heap.as_slice(),
893            &["a", "b", "c", "d", "e", "f", "g", "h", "i"]
894        );
895        assert!(heap.is_spilled());
896    }
897
898    #[test]
899    fn push_stack() {
900        let mut stack = SegmentsVec::from_slice(&["a", "b", "c"]);
901        stack.push("d");
902        stack.push("e");
903
904        assert_eq!(stack.as_slice(), &["a", "b", "c", "d", "e"]);
905        assert!(!stack.is_spilled());
906    }
907
908    #[test]
909    fn push_stack_spill() {
910        let mut stack = SegmentsVec::from_slice(&["a", "b", "c", "d", "e", "f", "g"]);
911        stack.push("h");
912
913        assert!(!stack.is_spilled());
914
915        stack.push("i");
916
917        assert_eq!(
918            stack.as_slice(),
919            &["a", "b", "c", "d", "e", "f", "g", "h", "i"]
920        );
921        assert!(stack.is_spilled());
922    }
923
924    #[test]
925    fn pop_stack() {
926        let mut stack = SegmentsVec::from_slice(&["a", "b", "c", "d", "e"]);
927        assert_eq!(stack.pop(), Some("e"));
928        assert_eq!(stack.pop(), Some("d"));
929        assert_eq!(stack.pop(), Some("c"));
930        assert_eq!(stack.pop(), Some("b"));
931        assert_eq!(stack.pop(), Some("a"));
932        assert_eq!(stack.pop(), None);
933
934        assert!(!stack.is_spilled());
935    }
936
937    #[test]
938    fn pop_heap() {
939        let mut heap = SegmentsVec::from_slice(&["a", "b", "c", "d", "e", "f", "g", "h", "i"]);
940
941        assert_eq!(heap.pop(), Some("i"));
942        assert_eq!(heap.pop(), Some("h"));
943        assert_eq!(heap.pop(), Some("g"));
944
945        assert!(heap.is_spilled());
946    }
947
948    #[test]
949    fn extend_from_slice_stack() {
950        let mut stack = SegmentsVec::from_slice(&["a", "b", "c"]);
951        stack.extend_from_slice(&["d", "e", "f"]);
952
953        assert_eq!(stack.as_slice(), &["a", "b", "c", "d", "e", "f"]);
954        assert!(!stack.is_spilled());
955    }
956
957    #[test]
958    fn extend_from_slice_stack_spill() {
959        let mut spilled = SegmentsVec::from_slice(&["a", "b", "c", "d", "e", "f"]);
960        spilled.extend_from_slice(&["g", "h", "i", "j"]);
961
962        assert_eq!(
963            spilled.as_slice(),
964            &["a", "b", "c", "d", "e", "f", "g", "h", "i", "j"]
965        );
966        assert!(spilled.is_spilled());
967    }
968
969    #[test]
970    fn extend_from_slice_heap() {
971        let mut heap = SegmentsVec::from_slice(&["a", "b", "c", "d", "e", "f", "g", "h", "i"]);
972        assert!(heap.is_spilled());
973
974        heap.extend_from_slice(&["j", "k", "l"]);
975
976        assert_eq!(
977            heap.as_slice(),
978            &["a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l"]
979        );
980    }
981
982    #[test]
983    fn extend_stack() {
984        let mut stack = SegmentsVec::from_slice(&["a", "b", "c"]);
985        stack.extend(["d", "e", "f"]);
986
987        assert_eq!(stack.as_slice(), &["a", "b", "c", "d", "e", "f"]);
988        assert!(!stack.is_spilled());
989    }
990
991    #[test]
992    fn extend_stack_spilled() {
993        let mut stack = SegmentsVec::from_slice(&["a", "b", "c", "d", "e", "f"]);
994        stack.extend(["g", "h", "i", "j"]);
995
996        assert_eq!(
997            stack.as_slice(),
998            &["a", "b", "c", "d", "e", "f", "g", "h", "i", "j"]
999        );
1000        assert!(stack.is_spilled());
1001    }
1002
1003    #[test]
1004    fn extend_heap() {
1005        let mut heap = SegmentsVec::from_slice(&["a", "b", "c", "d", "e", "f", "g", "h", "i"]);
1006        assert!(heap.is_spilled());
1007
1008        heap.extend(["j", "k", "l"]);
1009
1010        assert_eq!(
1011            heap.as_slice(),
1012            &["a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l"]
1013        );
1014    }
1015}
rustpython_ruff_python_ast/name.rs

rustpython_ruff_python_ast/
name.rs
