rib/type_inference/
inferred_expr.rs

1// Copyright 2024-2025 Golem Cloud
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::call_type::CallType;
16use crate::rib_compilation_error::RibCompilationError;
17use crate::{
18    DynamicParsedFunctionName, Expr, ExprVisitor, FunctionTypeRegistry, GlobalVariableTypeSpec,
19    RegistryKey,
20};
21use std::collections::HashSet;
22
23#[derive(Debug, Clone)]
24pub struct InferredExpr(Expr);
25
26impl InferredExpr {
27    pub fn get_expr(&self) -> &Expr {
28        &self.0
29    }
30
31    pub fn from_expr(
32        expr: Expr,
33        function_type_registry: &FunctionTypeRegistry,
34        type_spec: &Vec<GlobalVariableTypeSpec>,
35    ) -> Result<InferredExpr, RibCompilationError> {
36        let mut mutable_expr = expr;
37
38        mutable_expr.infer_types(function_type_registry, type_spec)?;
39
40        Ok(InferredExpr(mutable_expr))
41    }
42
43    // Only a fully inferred Rib can reliably tell us what are the exact
44    // function calls.
45    pub fn worker_invoke_calls(&self) -> Vec<DynamicParsedFunctionName> {
46        let mut expr = self.0.clone();
47        let mut worker_calls = vec![];
48        let mut visitor = ExprVisitor::bottom_up(&mut expr);
49
50        while let Some(expr) = visitor.pop_back() {
51            if let Expr::Call {
52                call_type: CallType::Function { function_name, .. },
53                ..
54            } = expr
55            {
56                worker_calls.push(function_name.clone());
57            }
58        }
59
60        worker_calls
61    }
62
63    pub fn worker_invoke_registry_keys(&self) -> HashSet<RegistryKey> {
64        let worker_calls = self.worker_invoke_calls();
65
66        let mut registry_keys = HashSet::new();
67
68        for call in worker_calls {
69            let keys = RegistryKey::registry_keys_of_function(&call);
70            registry_keys.extend(keys)
71        }
72
73        registry_keys
74    }
75}