starlark/values/types/string/
iter.rs

1/*
2 * Copyright 2018 The Starlark in Rust Authors.
3 * Copyright (c) Facebook, Inc. and its affiliates.
4 *
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
8 *
9 *     https://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 */
17
18//! Implementation of iterators for string type.
19
20use allocative::Allocative;
21use derive_more::Display;
22use starlark_derive::starlark_value;
23use starlark_derive::Freeze;
24use starlark_derive::NoSerialize;
25use starlark_derive::Trace;
26
27use crate as starlark;
28use crate::any::ProvidesStaticType;
29use crate::coerce::Coerce;
30use crate::values::typing::iter::StarlarkIter;
31use crate::values::FreezeResult;
32use crate::values::Heap;
33use crate::values::StarlarkValue;
34use crate::values::StringValue;
35use crate::values::StringValueLike;
36use crate::values::Value;
37use crate::values::ValueLike;
38use crate::values::ValueOfUnchecked;
39
40/// An opaque iterator over a string, produced by elems/codepoints
41#[derive(
42    Debug,
43    Trace,
44    Coerce,
45    Display,
46    Freeze,
47    NoSerialize,
48    ProvidesStaticType,
49    Allocative
50)]
51#[display("iterator")]
52#[repr(C)]
53struct StringIterableGen<'v, V: ValueLike<'v>> {
54    string: V::String,
55    produce_char: bool, // if not char, then int
56}
57
58pub(crate) fn iterate_chars<'v>(
59    string: StringValue<'v>,
60    heap: &'v Heap,
61) -> ValueOfUnchecked<'v, StarlarkIter<String>> {
62    ValueOfUnchecked::new(heap.alloc_complex(StringIterableGen::<'v, Value<'v>> {
63        string,
64        produce_char: true,
65    }))
66}
67
68pub(crate) fn iterate_codepoints<'v>(
69    string: StringValue<'v>,
70    heap: &'v Heap,
71) -> ValueOfUnchecked<'v, StarlarkIter<String>> {
72    ValueOfUnchecked::new(heap.alloc_complex(StringIterableGen::<'v, Value<'v>> {
73        string,
74        produce_char: false,
75    }))
76}
77
78#[starlark_value(type = "iterator")]
79impl<'v, V: ValueLike<'v>> StarlarkValue<'v> for StringIterableGen<'v, V>
80where
81    Self: ProvidesStaticType<'v>,
82{
83    unsafe fn iterate(&self, _me: Value<'v>, heap: &'v Heap) -> crate::Result<Value<'v>> {
84        // Lazy implementation: we allocate a tuple and then iterate over it.
85        let iter = if self.produce_char {
86            heap.alloc_tuple_iter(self.string.as_str().chars().map(|c| heap.alloc(c)))
87        } else {
88            heap.alloc_tuple_iter(
89                self.string
90                    .as_str()
91                    .chars()
92                    .map(|c| heap.alloc(u32::from(c))),
93            )
94        };
95        Ok(iter)
96    }
97}