transpile: Translate enums in switch as their own type instead of integers

c2rust-ast-builder/src/builder.rs

@@ -1233,6 +1233,20 @@ impl Builder {
         })
     }
 
+    pub fn tuple_struct_pat<Pa>(self, path: Pa, qself: Option<QSelf>, elems: Vec<Pat>) -> Pat
+    where
+        Pa: Make<Path>,
+    {
+        let path = path.make(&self);
+        Pat::TupleStruct(PatTupleStruct {
+            attrs: self.attrs,
+            qself,
+            path,
+            paren_token: token::Paren(self.span),
+            elems: punct(elems),
+        })
+    }
+
     // Types
 
     pub fn barefn_ty(self, decl: BareFnTyParts) -> Box<Type> {

c2rust-transpile/src/cfg/mod.rs

@@ -34,7 +34,7 @@ use std::ops::Index;
 use std::rc::Rc;
 use std::{fmt, io};
 use syn::Lit;
-use syn::{spanned::Spanned, Arm, Expr, Pat, Stmt};
+use syn::{punctuated::Punctuated, spanned::Spanned, Arm, Expr, Pat, Stmt};
 
 use failure::format_err;
 use indexmap::indexset;
@@ -448,6 +448,7 @@ impl GenTerminator<StructureLabel<StmtOrDecl>> {
 pub struct SwitchCases {
     cases: Vec<(Pat, Label)>,
     default: Option<Label>,
+    override_type_id: Option<CQualTypeId>,
 }
 
 /// A Rust statement, or a C declaration, or a comment
@@ -1934,42 +1935,45 @@ impl CfgBuilder {
                 self.add_wip_block(wip, Jump(this_label.clone()));
 
                 // Case
-                let resolved = translator.ast_context.unwrap_cast_expr(case_expr);
-                let branch = match translator.ast_context.index_unwrap_parens(resolved).kind {
-                    CExprKind::Literal(..) | CExprKind::ConstantExpr(_, _, Some(_)) => {
-                        match translator
-                            .convert_expr(ctx.used(), resolved, None)?
-                            .to_pure_expr()
+
+                let switch_case = self.switch_expr_cases.last_mut().ok_or_else(|| {
+                    format_err!(
+                        "Cannot find the 'switch' wrapping this ({:?}) 'case' statement",
+                        stmt_id,
+                    )
+                })?;
+
+                let expr = translator
+                    .convert_expr(
+                        ctx.const_().pattern().used(),
+                        case_expr,
+                        switch_case.override_type_id,
+                    )
+                    .ok()
+                    .and_then(WithStmts::to_pure_expr);
+                let pat = expr.and_then(|expr| expr_to_pat(*expr)).unwrap_or_else(|| {
+                    let mut pat = match cie {
+                        ConstIntExpr::U(n) => mk().lit_pat(mk().int_unsuffixed_lit(n)),
+                        ConstIntExpr::I(n) => mk().lit_pat(mk().int_unsuffixed_lit(n)),
+                    };
+
+                    if let Some(override_ty) = switch_case.override_type_id {
+                        if let CTypeKind::Enum(enum_id) =
+                            translator.ast_context.resolve_type(override_ty.ctype).kind
                         {
-                            Some(expr) => match *expr {
-                                Expr::Lit(lit) => Some(mk().lit_pat(lit.lit)),
-                                Expr::Path(path) => Some(mk().path_pat(path.path, path.qself)),
-                                _ => None,
-                            },
-                            _ => None,
+                            let enum_name = translator
+                                .type_converter
+                                .borrow()
+                                .resolve_decl_name(enum_id)
+                                .unwrap();
+                            pat = mk().tuple_struct_pat(enum_name.as_str(), None, vec![pat]);
                         }
                     }
-                    _ => None,
-                };
 
-                let pat = match branch {
-                    Some(pat) => pat,
-                    None => match cie {
-                        ConstIntExpr::U(n) => mk().lit_pat(mk().int_unsuffixed_lit(n)),
-                        ConstIntExpr::I(n) => mk().lit_pat(mk().int_unsuffixed_lit(n)),
-                    },
-                };
+                    pat
+                });
 
-                self.switch_expr_cases
-                    .last_mut()
-                    .ok_or_else(|| {
-                        format_err!(
-                            "Cannot find the 'switch' wrapping this ({:?}) 'case' statement",
-                            stmt_id,
-                        )
-                    })?
-                    .cases
-                    .push((pat, this_label.clone()));
+                switch_case.cases.push((pat, this_label.clone()));
 
                 // Sub stmt
                 let sub_stmt_next = self.convert_stmt_help(
@@ -2015,8 +2019,30 @@ impl CfgBuilder {
                 let body_label = self.fresh_label();
 
                 // Convert the condition
+
+                let mut override_type_id = None;
+
+                // If the condition is an implicit cast from an enum to its integral type,
+                // override the type to that of the enum.
+                if let CExprKind::ImplicitCast(target_type_id, castee_id, ..) =
+                    translator.ast_context[scrutinee].kind
+                {
+                    let castee_kind = &translator.ast_context[castee_id].kind;
+                    let castee_type_id = castee_kind.get_qual_type().unwrap();
+                    let castee_type_kind = &translator
+                        .ast_context
+                        .resolve_type(castee_type_id.ctype)
+                        .kind;
+
+                    if let &CTypeKind::Enum(enum_id) = castee_type_kind {
+                        if target_type_id == translator.enum_integral_type(enum_id) {
+                            override_type_id = Some(castee_type_id);
+                        }
+                    }
+                }
+
                 let (stmts, val) = translator
-                    .convert_expr(ctx.used(), scrutinee, None)?
+                    .convert_expr(ctx.used(), scrutinee, override_type_id)?
                     .discard_unsafe();
                 wip.extend(stmts);
 
@@ -2028,7 +2054,10 @@ impl CfgBuilder {
                 let saw_unmatched_case = self.last_per_stmt_mut().saw_unmatched_case;
                 let saw_unmatched_default = self.last_per_stmt_mut().saw_unmatched_default;
                 self.break_labels.push(next_label.clone());
-                self.switch_expr_cases.push(SwitchCases::default());
+                self.switch_expr_cases.push(SwitchCases {
+                    override_type_id,
+                    ..Default::default()
+                });
 
                 let body_stuff = self.convert_stmt_help(
                     translator,
@@ -2415,3 +2444,173 @@ impl Cfg<Label, StmtOrDecl> {
         Ok(())
     }
 }
+
+fn expr_to_pat(expr: Expr) -> Option<Pat> {
+    use syn::{
+        ExprArray, ExprCall, ExprLit, ExprParen, ExprPath, ExprReference, ExprStruct, ExprTuple,
+        ExprUnary, FieldPat, FieldValue, LitInt, PatLit, PatParen, PatPath, PatReference, PatSlice,
+        PatStruct, PatTuple, PatTupleStruct, UnOp,
+    };
+
+    match expr {
+        Expr::Array(ExprArray {
+            attrs,
+            bracket_token,
+            elems,
+        }) => {
+            let elems = punctuated_expr_to_pat(elems)?;
+
+            Some(Pat::Slice(PatSlice {
+                attrs,
+                bracket_token,
+                elems,
+            }))
+        }
+
+        Expr::Call(ExprCall {
+            attrs,
+            func,
+            paren_token,
+            args,
+        }) => {
+            let (qself, path) = match *func {
+                Expr::Path(ExprPath { qself, path, .. }) => (qself, path),
+                _ => return None,
+            };
+            let elems = punctuated_expr_to_pat(args)?;
+
+            Some(Pat::TupleStruct(PatTupleStruct {
+                attrs,
+                qself,
+                path,
+                paren_token,
+                elems,
+            }))
+        }
+
+        Expr::Lit(ExprLit { attrs, lit }) => Some(Pat::Lit(PatLit { attrs, lit })),
+
+        Expr::Paren(ExprParen {
+            attrs,
+            paren_token,
+            expr,
+        }) => {
+            let pat = Box::new(expr_to_pat(*expr)?);
+            Some(Pat::Paren(PatParen {
+                attrs,
+                paren_token,
+                pat,
+            }))
+        }
+
+        Expr::Path(ExprPath { attrs, qself, path }) => {
+            Some(Pat::Path(PatPath { attrs, qself, path }))
+        }
+
+        Expr::Range(range) => Some(Pat::Range(range)),
+
+        Expr::Reference(ExprReference {
+            attrs,
+            and_token,
+            mutability,
+            expr,
+        }) => {
+            let pat = Box::new(expr_to_pat(*expr)?);
+
+            Some(Pat::Reference(PatReference {
+                attrs,
+                and_token,
+                mutability,
+                pat,
+            }))
+        }
+
+        Expr::Struct(ExprStruct {
+            attrs,
+            qself,
+            path,
+            brace_token,
+            fields,
+            dot2_token: None,
+            rest: None,
+        }) => {
+            let fields = fields
+                .into_iter()
+                .map(|field| {
+                    let FieldValue {
+                        attrs,
+                        member,
+                        colon_token,
+                        expr,
+                    } = field;
+                    let pat = Box::new(expr_to_pat(expr)?);
+
+                    Some(FieldPat {
+                        attrs,
+                        member,
+                        colon_token,
+                        pat,
+                    })
+                })
+                .collect::<Option<_>>()?;
+
+            Some(Pat::Struct(PatStruct {
+                attrs,
+                qself,
+                path,
+                brace_token,
+                fields,
+                rest: None,
+            }))
+        }
+
+        Expr::Tuple(ExprTuple {
+            attrs,
+            paren_token,
+            elems,
+        }) => {
+            let elems = punctuated_expr_to_pat(elems)?;
+
+            Some(Pat::Tuple(PatTuple {
+                attrs,
+                paren_token,
+                elems,
+            }))
+        }
+
+        // There is no equivalent `PatUnary`, but the negative sign can be folded into the literal.
+        Expr::Unary(ExprUnary {
+            attrs,
+            op: UnOp::Neg(_),
+            expr,
+        }) => {
+            let Expr::Lit(ExprLit {
+                attrs: _,
+                lit: Lit::Int(lit_int),
+            }) = *expr else {
+                return None
+            };
+
+            let repr = format!("-{}{}", lit_int.base10_digits(), lit_int.suffix());
+            let lit = Lit::Int(LitInt::new(&repr, lit_int.span()));
+            Some(Pat::Lit(PatLit { attrs, lit }))
+        }
+
+        _ => None,
+    }
+}
+
+fn punctuated_expr_to_pat(
+    elems: Punctuated<Expr, syn::Token![,]>,
+) -> Option<Punctuated<Pat, syn::Token![,]>> {
+    use syn::punctuated::Pair;
+
+    elems
+        .into_pairs()
+        .map(|pair| {
+            let (expr, token) = pair.into_tuple();
+            let pat = expr_to_pat(expr)?;
+            Some(Pair::new(pat, token))
+        })
+        .collect()
+}