statement.rs

// Copyright 2026 Oxide Computer Company

use crate::{
    expression::{slice_mask, ExpressionGenerator},
    is_header, is_header_member, is_rust_reference, rust_type,
};
use p4::ast::{
    Call, Control, DeclarationInfo, Direction, Expression, ExpressionKind,
    NameInfo, Parser, Statement, StatementBlock, Transition, Type, AST,
};
use p4::hlir::Hlir;
use proc_macro2::TokenStream;
use quote::{format_ident, quote};
use std::collections::HashMap;

#[derive(Debug)]
pub(crate) enum StatementContext<'a> {
    Control(&'a Control),
    #[allow(dead_code)]
    Parser(&'a Parser),
}

pub(crate) struct StatementGenerator<'a> {
    hlir: &'a Hlir,
    ast: &'a AST,
    context: StatementContext<'a>,
}

impl<'a> StatementGenerator<'a> {
    pub fn new(
        ast: &'a AST,
        hlir: &'a Hlir,
        context: StatementContext<'a>,
    ) -> Self {
        Self { ast, hlir, context }
    }

    pub(crate) fn generate_block(
        &self,
        sb: &StatementBlock,
        names: &mut HashMap<String, NameInfo>,
    ) -> TokenStream {
        let mut ts = TokenStream::new();
        for stmt in &sb.statements {
            ts.extend(self.generate_statement(stmt, names));
        }
        ts
    }

    pub(crate) fn generate_statement(
        &self,
        stmt: &Statement,
        names: &mut HashMap<String, NameInfo>,
    ) -> TokenStream {
        match stmt {
            Statement::Empty => TokenStream::new(),
            Statement::Assignment(lval, xpr) => {
                let eg = ExpressionGenerator::new(self.hlir);

                let lhs = eg.generate_lvalue(lval);

                let rhs = eg.generate_expression(xpr.as_ref());
                let rhs_ty = self
                    .hlir
                    .expression_types
                    .get(xpr.as_ref())
                    .unwrap_or_else(|| {
                        panic!("codegen type not found for {:#?}", xpr)
                    });

                let name_info =
                    self.hlir.lvalue_decls.get(lval).unwrap_or_else(|| {
                        panic!("codegen name not resolved for {:#?}", lval)
                    });

                let converted = rhs_ty != &name_info.ty;
                let rhs = if converted {
                    let converter = self.converter(rhs_ty, &name_info.ty);
                    let rhs = if matches!(rhs_ty, Type::Bit(_)) {
                        quote! { (#rhs).to_owned() }
                    } else {
                        rhs
                    };
                    quote!( #converter(#rhs) )
                } else {
                    rhs
                };

                if is_header_member(lval, self.hlir) {
                    return if converted {
                        quote! { #lhs = #rhs; }
                    } else {
                        quote! { #lhs = #rhs.clone(); }
                    };
                }

                let rhs = if converted {
                    rhs
                } else if let Type::Bit(_) = rhs_ty {
                    // TODO eww, to better to figure out precisely when to_owned
                    // and clone are needed
                    quote! { #rhs.to_owned() }
                } else if let Type::UserDefined(_) = rhs_ty {
                    quote! { #rhs.clone() }
                } else {
                    rhs
                };

                if is_rust_reference(lval, names) {
                    quote! { *#lhs = #rhs; }
                } else {
                    quote! { #lhs = #rhs; }
                }
            }
            Statement::SliceAssignment(lval, hi, lo, xpr) => {
                let eg = ExpressionGenerator::new(self.hlir);
                let lhs = eg.generate_lvalue(lval);
                let mut rhs = eg.generate_expression(xpr.as_ref());

                let name_info =
                    self.hlir.lvalue_decls.get(lval).unwrap_or_else(|| {
                        panic!(
                            "unresolved lvalue {lval:#?} in slice assignment"
                        )
                    });
                let field_width = match &name_info.ty {
                    Type::Bit(w) => *w,
                    ty => panic!(
                        "slice assignment on non-bit type {ty:?} reached codegen"
                    ),
                };

                let (hi_val, lo_val) =
                    ExpressionGenerator::slice_bounds(hi, lo);

                let slice_ty = Type::Bit(hi_val - lo_val + 1);
                let rhs_ty = self
                    .hlir
                    .expression_types
                    .get(xpr.as_ref())
                    .unwrap_or_else(|| panic!("type for {xpr:#?}"));

                let converted = rhs_ty != &slice_ty;
                if converted {
                    let converter = self.converter(rhs_ty, &slice_ty);
                    if matches!(rhs_ty, Type::Bit(_)) {
                        rhs = quote! { (#rhs).to_owned() };
                    }
                    rhs = quote! { #converter(#rhs) };
                }

                let mask_val = slice_mask(hi_val - lo_val + 1);
                quote! {
                    {
                        let __rhs_val: u128 = #rhs.load_le();
                        if #lhs.is_empty() {
                            #lhs.resize(#field_width, false);
                        }
                        let __lhs_val: u128 = #lhs.load_le();
                        let __mask: u128 = #mask_val << #lo_val;
                        let __new = (__lhs_val & !__mask)
                            | ((__rhs_val & #mask_val) << #lo_val);
                        #lhs.store_le(__new);
                    }
                }
            }
            Statement::Call(c) => match &self.context {
                StatementContext::Control(control) => {
                    let mut ts = TokenStream::new();
                    self.generate_control_body_call(control, c, &mut ts);
                    ts
                }
                StatementContext::Parser(parser) => {
                    let mut ts = TokenStream::new();
                    self.generate_parser_body_call(parser, c, &mut ts);
                    ts
                }
            },
            Statement::If(ifb) => {
                let eg = ExpressionGenerator::new(self.hlir);
                let predicate = eg.generate_expression(ifb.predicate.as_ref());
                let block = self.generate_block(&ifb.block, names);
                let mut ts = quote! {
                    if #predicate { #block }
                };
                for ei in &ifb.else_ifs {
                    let predicate =
                        eg.generate_expression(ei.predicate.as_ref());
                    let block = self.generate_block(&ei.block, names);
                    ts.extend(quote! {else if #predicate { #block }})
                }
                if let Some(eb) = &ifb.else_block {
                    let block = self.generate_block(eb, names);
                    ts.extend(quote! {else { #block }})
                }
                ts
            }
            Statement::Variable(v) => {
                let name = format_ident!("{}", v.name);
                let ty = rust_type(&v.ty);
                let initializer = match &v.initializer {
                    Some(xpr) => self.initializer(xpr, &v.ty),
                    None => match &v.ty {
                        Type::Bit(w) => quote! { bitvec![u8, Msb0; 0; #w] },
                        _ => quote! { #ty::default() },
                    },
                };
                names.insert(
                    v.name.clone(),
                    NameInfo {
                        ty: v.ty.clone(),
                        decl: DeclarationInfo::Local,
                    },
                );

                //TODO determine for real
                let needs_mut = true;

                if needs_mut {
                    quote! {
                        let mut #name: #ty = #initializer;
                    }
                } else {
                    quote! {
                        let #name: #ty = #initializer;
                    }
                }
            }
            Statement::Constant(c) => {
                let name = format_ident!("{}", c.name);
                let ty = rust_type(&c.ty);
                let initializer = self.initializer(&c.initializer, &c.ty);
                quote! {
                    let #name: #ty = #initializer;
                }
            }
            Statement::Transition(transition) => {
                let parser = match self.context {
                    StatementContext::Parser(p) => p,
                    _ => {
                        panic!(
                            "transition statement outside parser: {:#?}",
                            transition,
                        )
                    }
                };
                match transition {
                    Transition::Reference(next_state) => {
                        match next_state.name.as_str() {
                            "accept" => quote! { return true; },
                            "reject" => quote! { return false; },
                            state_ref => {
                                let state_name = format_ident!(
                                    "{}_{}",
                                    parser.name,
                                    state_ref
                                );
                                let mut args = Vec::new();
                                for arg in &parser.parameters {
                                    let name = format_ident!("{}", arg.name);
                                    args.push(quote! { #name });
                                }
                                quote! {
                                    softnpu_provider::parser_transition!(||(#state_ref));
                                    return #state_name( #(#args),* );
                                }
                            }
                        }
                    }
                    Transition::Select(_) => {
                        todo!();
                    }
                }
            }
            Statement::Return(xpr) => {
                let eg = ExpressionGenerator::new(self.hlir);
                if let Some(xpr) = xpr {
                    let xp = eg.generate_expression(xpr.as_ref());
                    quote! { return #xp; }
                } else {
                    quote! { return }
                }
            }
        }
    }

    fn generate_parser_body_call(
        &self,
        parser: &Parser,
        c: &Call,
        tokens: &mut TokenStream,
    ) {
        let lval: Vec<TokenStream> = c
            .lval
            .name
            .split('.')
            .map(|x| format_ident!("{}", x))
            .map(|x| quote! { #x })
            .collect();

        let mut args = Vec::new();
        for a in &c.args {
            match &a.kind {
                ExpressionKind::Lvalue(lvarg) => {
                    let parts: Vec<&str> = lvarg.name.split('.').collect();
                    let root = parts[0];
                    let mut mut_arg = false;
                    for parg in &parser.parameters {
                        if parg.name == root {
                            match parg.direction {
                                Direction::Out | Direction::InOut => {
                                    mut_arg = true;
                                }
                                _ => {}
                            }
                        }
                    }
                    let lvref: Vec<TokenStream> = parts
                        .iter()
                        .map(|x| format_ident!("{}", x))
                        .map(|x| quote! { #x })
                        .collect();
                    if mut_arg {
                        args.push(quote! { &mut #(#lvref).* });
                    } else {
                        args.push(quote! { #(#lvref).* });
                    }
                }
                x => todo!("extern arg {:?}", x),
            }
        }
        tokens.extend(quote! {
            #(#lval).* ( #(#args),* );
        });
    }

    fn generate_control_body_call(
        &self,
        control: &Control,
        c: &Call,
        tokens: &mut TokenStream,
    ) {
        if c.lval.name.split('.').count() < 2 {
            self.generate_action_call(control, c, tokens);
            return;
        }

        match c.lval.leaf() {
            "apply" => {
                self.generate_control_apply_body_call(control, c, tokens);
            }
            "setValid" => {
                self.generate_header_set_validity(c, tokens, true);
            }
            "setInvalid" => {
                self.generate_header_set_validity(c, tokens, false);
            }
            "isValid" => {
                self.generate_header_get_validity(c, tokens);
            }
            _ => {
                // assume we are at an extern call

                // TODO check the extern call against defined externs in checker
                // before we get here

                self.generate_control_extern_call(control, c, tokens);
            }
        }
    }

    fn generate_action_call(
        &self,
        control: &Control,
        c: &Call,
        tokens: &mut TokenStream,
    ) {
        let eg = ExpressionGenerator::new(self.hlir);
        let mut args = Vec::new();

        // Control parameters come first in the action function signature
        // (see `generate_control_action` in control.rs), followed by
        // extern references, and then action-specific parameters.
        for a in &control.parameters {
            let arg = format_ident!("{}", a.name);
            args.push(quote! { #arg });
        }

        // Pass externs instantiated at control scope to actions.
        for x in &control.variables {
            if let Type::UserDefined(typename) = &x.ty {
                if self.ast.get_extern(typename).is_some() {
                    let arg = format_ident!("{}", x.name);
                    args.push(quote! { & #arg });
                }
            }
        }

        // Action-specific arguments come last. We clone lvalue args to avoid
        // moving out from mutable references.
        for a in &c.args {
            let arg_xpr = eg.generate_expression(a.as_ref());
            if matches!(a.kind, ExpressionKind::Lvalue(_)) {
                args.push(quote! { #arg_xpr.clone() });
            } else {
                args.push(arg_xpr);
            }
        }

        let lvref: Vec<TokenStream> = c
            .lval
            .name
            .split('.')
            .map(|x| format_ident!("{}_action_{x}", control.name))
            .map(|x| quote! { #x })
            .collect();

        tokens.extend(quote! {
            #(#lvref).*(#(#args),*);
        })
    }

    fn generate_control_extern_call(
        &self,
        _control: &Control,
        c: &Call,
        tokens: &mut TokenStream,
    ) {
        let eg = ExpressionGenerator::new(self.hlir);
        let mut args = Vec::new();

        for a in &c.args {
            let arg_xpr = eg.generate_expression(a.as_ref());
            args.push(arg_xpr);
        }

        let lvref: Vec<TokenStream> = c
            .lval
            .name
            .split('.')
            .map(|x| format_ident!("{}", x))
            .map(|x| quote! { #x })
            .collect();

        tokens.extend(quote! {
            #(#lvref).*(#(#args),*);
        })
    }

    fn generate_control_apply_body_call(
        &self,
        control: &Control,
        c: &Call,
        tokens: &mut TokenStream,
    ) {
        let name_info = self
            .hlir
            .lvalue_decls
            .get(&c.lval.pop_right())
            .unwrap_or_else(|| {
                panic!("codegen: lval root for {:#?} not found in hlir", c.lval)
            });

        let control_instance = match &name_info.ty {
            Type::UserDefined(name) => {
                self.ast.get_control(name).unwrap_or_else(|| {
                    panic!("codegen: control {} not found in AST", name,)
                })
            }
            Type::Table => control,
            t => panic!(
                "call {:#?} references non-user-defined type {:#?}",
                c, t
            ),
        };

        let root = c.lval.root();

        // This is a call to another control instance
        if control_instance.name != control.name {
            let eg = ExpressionGenerator::new(self.hlir);
            let mut locals = Vec::new();
            let mut args = Vec::new();
            for (i, a) in c.args.iter().enumerate() {
                let arg_xpr = eg.generate_expression(a.as_ref());
                let et = self.hlir.expression_types.get(a.as_ref()).unwrap();
                let local_name = format_ident!("arg{}", i);
                let param_ty = &control_instance.parameters[i].ty;

                if et != param_ty {
                    let converter = self.converter(et, param_ty);
                    let arg_xpr = if matches!(et, Type::Bit(_)) {
                        quote! { (#arg_xpr).to_owned() }
                    } else {
                        arg_xpr
                    };
                    locals.push(quote! {
                        let mut #local_name = #converter(#arg_xpr);
                    });
                    args.push(quote! { &#local_name });
                    continue;
                }

                if let ExpressionKind::BitLit(_, _) = a.as_ref().kind {
                    locals.push(quote! {
                        let mut #local_name = #arg_xpr;
                    })
                }

                match control_instance.parameters[i].direction {
                    Direction::Out | Direction::InOut => {
                        match &a.as_ref().kind {
                            ExpressionKind::Lvalue(lval) => {
                                let name_info = self
                                    .hlir
                                    .lvalue_decls
                                    .get(lval)
                                    .unwrap_or_else(|| {
                                        panic!(
                                        "codegen: lvalue resolve fail {:#?}",
                                        lval
                                    )
                                    });
                                match name_info.decl {
                                    // if this is a control parameter, it should
                                    // already have the correct reference
                                    // annotations
                                    DeclarationInfo::Parameter(_) => {
                                        args.push(arg_xpr);
                                    }
                                    _ => {
                                        args.push(quote! {&mut #arg_xpr});
                                    }
                                }
                            }
                            ExpressionKind::BitLit(_, _) => {
                                args.push(quote! { &mut #local_name })
                            }
                            _ => {
                                args.push(quote! {&mut #arg_xpr});
                            }
                        }
                    }
                    _ => match &a.as_ref().kind {
                        ExpressionKind::BitLit(_, _) => {
                            args.push(quote! { &#local_name })
                        }
                        // if the argument is an lvalue and a rust type that
                        // gets passed by reference, take a ref
                        ExpressionKind::Lvalue(_) => match et {
                            Type::Bit(_) | Type::Varbit(_) | Type::Int(_) => {
                                args.push(quote! { &#arg_xpr });
                            }
                            _ => {
                                args.push(arg_xpr);
                            }
                        },
                        _ => {
                            args.push(arg_xpr);
                        }
                    },
                }
            }

            let tables = control_instance.tables(self.ast);
            for (cs, table) in tables {
                let qtn =
                    crate::qualified_table_function_name(None, &cs, table);
                let name = format_ident!("{}_{}", c.lval.root(), qtn);
                args.push(quote! { #name });
            }

            let cname = &control_instance.name;
            let call = format_ident!("{}_apply", control_instance.name);

            tokens.extend(quote! {
                softnpu_provider::control_apply!(||(#cname));
                #(#locals);*
                #call(#(#args),*);
            });

            return;
        }

        // this is a local table

        let table = match control_instance.get_table(root) {
            Some(table) => table,
            None => {
                panic!(
                    "codegen: table {} not found in control {} decl: {:#?}",
                    root, control.name, name_info,
                );
            }
        };

        //
        // match an action based on the key material
        //

        let table_name = format_ident!("{}", c.lval.root());

        let table_name_str =
            format!("{}_table_{}", control_instance.name, table.name,);

        let mut action_args = Vec::new();
        for p in &control.parameters {
            let name = format_ident!("{}", p.name);
            action_args.push(quote! { #name });
        }

        for var in &control.variables {
            let name = format_ident!("{}", var.name);
            if let Type::UserDefined(typename) = &var.ty {
                if self.ast.get_extern(typename).is_some() {
                    action_args.push(quote! { &#name });
                }
            }
        }

        let mut selector_components = Vec::new();
        for (lval, _match_kind) in &table.key {
            let lvref: Vec<TokenStream> = lval
                .name
                .split('.')
                .map(|x| format_ident!("{}", x))
                .map(|x| quote! { #x })
                .collect();

            // determine if this lvalue references a header or a struct,
            // if it's a header there's a bit of extra unsrapping we
            // need to do to match the selector against the value.

            let names = control.names();

            if lval.degree() > 1
                && is_header(&lval.pop_right(), self.ast, &names)
            {
                //TODO: to_biguint is bad here, copying on data path
                selector_components.push(quote! {
                    p4rs::bitvec_to_biguint(
                        &#(#lvref).*
                    ).value
                });
            } else {
                selector_components.push(quote! {
                    p4rs::bitvec_to_biguint(&#(#lvref).*).value
                });
            }
        }
        let default_action =
            format_ident!("{}_action_{}", control.name, table.default_action);
        tokens.extend(quote! {
            let matches = #table_name.match_selector(
                &[#(#selector_components),*]
            );
            if matches.len() > 0 {
                softnpu_provider::control_table_hit!(||#table_name_str);
                (matches[0].action)(#(#action_args),*)
            }
        });
        if table.default_action != "NoAction" {
            tokens.extend(quote! {
                else {
                    softnpu_provider::control_table_miss!(||#table_name_str);
                    #default_action(#(#action_args),*);
                }
            });
        } else {
            tokens.extend(quote! {
                else {
                    softnpu_provider::control_table_miss!(||#table_name_str);
                }
            });
        }
    }

    fn generate_header_set_validity(
        &self,
        c: &Call,
        tokens: &mut TokenStream,
        valid: bool,
    ) {
        let lhs: Vec<TokenStream> = c
            .lval
            .pop_right()
            .name
            .split('.')
            .map(|x| format_ident!("{}", x))
            .map(|x| quote! { #x })
            .collect();
        if valid {
            tokens.extend(quote! {
                #(#lhs).*.set_valid();
            });
        } else {
            tokens.extend(quote! {
                #(#lhs).*.set_invalid();
            });
        }
    }

    fn generate_header_get_validity(&self, c: &Call, tokens: &mut TokenStream) {
        let lhs: Vec<TokenStream> = c
            .lval
            .pop_right()
            .name
            .split('.')
            .map(|x| format_ident!("{}", x))
            .map(|x| quote! { #x })
            .collect();
        tokens.extend(quote! {
            #(#lhs).*.is_valid()
        });
    }

    /// Generate the value bound by a declaration like `bit<W> x = e;` or
    /// `const bit<W> c = e;`.
    ///
    /// Lvalues are cloned, and the `converter` fn applies when the expression's
    /// type differs from the declared one.
    fn initializer(&self, xpr: &Expression, ty: &Type) -> TokenStream {
        let eg = ExpressionGenerator::new(self.hlir);
        let mut ini = eg.generate_expression(xpr);
        if let ExpressionKind::Lvalue(_) = xpr.kind {
            ini = quote! { #ini.clone() };
        }

        let ini_ty = self
            .hlir
            .expression_types
            .get(xpr)
            .unwrap_or_else(|| panic!("type for expression {xpr:#?}"));

        if ini_ty != ty {
            let converter = self.converter(ini_ty, ty);
            quote! { #converter(#ini) }
        } else {
            ini
        }
    }

    /// Generate a closure converting a value of type `from` to `to`.
    ///
    /// Integer literals are truncated to `bit<W>` as specified in P4-16
    /// 8.11.1. Implicit bit widening is an x4c extension implemented through
    /// `p4rs::bitvec_resize`.
    fn converter(&self, from: &Type, to: &Type) -> TokenStream {
        match (from, to) {
            (Type::Int(_), Type::Bit(0)) => {
                quote! { (|_: i128| BitVec::<u8, Msb0>::new()) }
            }
            (Type::Int(_), Type::Bit(w)) => {
                quote! {
                    (|__v: i128| {
                        let mut __x = bitvec![u8, Msb0; 0; #w];
                        __x.store_le(__v as u128);
                        __x
                    })
                }
            }
            (Type::Bit(_), Type::Bit(w)) => {
                quote! { (|__bv| p4rs::bitvec_resize(__bv, #w)) }
            }
            _ => unreachable!("hlir rejects converting {from} to {to}"),
        }
    }
}