Skip to main content

luaur_analysis/functions/
keyof_function_impl.rs

1use crate::enums::reduction::Reduction;
2use crate::functions::compute_keys_of::compute_keys_of;
3use crate::functions::follow_type::follow_type_id;
4use crate::records::singleton_type::SingletonType;
5use crate::records::string_singleton::StringSingleton;
6use crate::records::type_function_context::TypeFunctionContext;
7use crate::records::type_function_reduction_result::TypeFunctionReductionResult;
8use crate::records::union_type::UnionType;
9use crate::type_aliases::singleton_variant::SingletonVariant;
10use crate::type_aliases::type_id::TypeId;
11use crate::type_aliases::type_pack_id::TypePackId;
12use alloc::collections::BTreeSet;
13use alloc::vec::Vec;
14use luaur_common::macros::luau_assert::LUAU_ASSERT;
15use luaur_common::records::dense_hash_set::DenseHashSet;
16
17pub fn keyof_function_impl(
18    type_params: Vec<TypeId>,
19    pack_params: Vec<TypePackId>,
20    ctx: *mut TypeFunctionContext,
21    is_raw: bool,
22) -> TypeFunctionReductionResult {
23    let ctx_ref = unsafe { &*ctx };
24    if type_params.len() != 1 || !pack_params.is_empty() {
25        unsafe {
26            (*ctx_ref.ice.as_ptr()).ice_string(
27                "keyof type function: encountered a type function instance without the required argument structure",
28            )
29        };
30        LUAU_ASSERT!(false);
31    }
32
33    let make_result = |result, reduction_status| TypeFunctionReductionResult {
34        result,
35        reduction_status,
36        blocked_types: Vec::new(),
37        blocked_packs: Vec::new(),
38        error: None,
39        messages: Vec::new(),
40    };
41
42    let operand_ty = unsafe { follow_type_id(type_params[0]) };
43    let Some(norm_ty) = (unsafe { (*ctx_ref.normalizer.as_ptr()).try_normalize(operand_ty) })
44    else {
45        return make_result(None, Reduction::MaybeOk);
46    };
47
48    if norm_ty.has_tables() == norm_ty.has_extern_types() {
49        return make_result(None, Reduction::Erroneous);
50    }
51
52    if norm_ty.has_tops()
53        || norm_ty.has_booleans()
54        || norm_ty.has_errors()
55        || norm_ty.has_nils()
56        || norm_ty.has_numbers()
57        || norm_ty.has_strings()
58        || norm_ty.has_threads()
59        || norm_ty.has_buffers()
60        || norm_ty.has_functions()
61        || norm_ty.has_tyvars()
62    {
63        return make_result(None, Reduction::Erroneous);
64    }
65
66    let mut keys = BTreeSet::new();
67
68    if norm_ty.has_extern_types() {
69        LUAU_ASSERT!(!norm_ty.has_tables());
70        let mut seen: DenseHashSet<TypeId> = DenseHashSet::new(core::ptr::null());
71
72        let mut extern_types = norm_ty.extern_types.ordering.iter().copied();
73        let Some(first) = extern_types.next() else {
74            return make_result(None, Reduction::Erroneous);
75        };
76
77        if !compute_keys_of(first, &mut keys, &mut seen, is_raw, ctx) {
78            return make_result(
79                Some(unsafe { ctx_ref.builtins.as_ref().stringType }),
80                Reduction::MaybeOk,
81            );
82        }
83
84        for extern_ty in extern_types {
85            seen.clear();
86            let mut local_keys = BTreeSet::new();
87            if compute_keys_of(extern_ty, &mut local_keys, &mut seen, is_raw, ctx) {
88                keys.retain(|key| local_keys.contains(key));
89            }
90        }
91    }
92
93    if norm_ty.has_tables() {
94        LUAU_ASSERT!(!norm_ty.has_extern_types());
95        let mut seen: DenseHashSet<TypeId> = DenseHashSet::new(core::ptr::null());
96
97        let mut tables = norm_ty.tables.order.iter().copied();
98        let Some(first) = tables.next() else {
99            return make_result(None, Reduction::Erroneous);
100        };
101
102        if !compute_keys_of(first, &mut keys, &mut seen, is_raw, ctx) {
103            return make_result(
104                Some(unsafe { ctx_ref.builtins.as_ref().stringType }),
105                Reduction::MaybeOk,
106            );
107        }
108
109        for table in tables {
110            seen.clear();
111            let mut local_keys = BTreeSet::new();
112            if compute_keys_of(table, &mut local_keys, &mut seen, is_raw, ctx) {
113                keys.retain(|key| local_keys.contains(key));
114            }
115        }
116    }
117
118    if keys.is_empty() {
119        return make_result(
120            Some(unsafe { ctx_ref.builtins.as_ref().neverType }),
121            Reduction::MaybeOk,
122        );
123    }
124
125    let mut singletons = Vec::new();
126    for key in keys {
127        singletons.push(unsafe {
128            (*ctx_ref.arena.as_ptr()).add_type(SingletonType::singleton_type(SingletonVariant::V1(
129                StringSingleton::new(key),
130            )))
131        });
132    }
133
134    if singletons.len() == 1 {
135        return make_result(Some(singletons[0]), Reduction::MaybeOk);
136    }
137
138    make_result(
139        Some(unsafe {
140            (*ctx_ref.arena.as_ptr()).add_type(UnionType {
141                options: singletons,
142            })
143        }),
144        Reduction::MaybeOk,
145    )
146}