From 375e94269dfe9cc2a82f148e5dab8f231180b6be Mon Sep 17 00:00:00 2001
From: Stuart Pernsteiner <spernsteiner@galois.com>
Date: Wed, 10 Jul 2024 16:20:28 -0700
Subject: [PATCH] analyze: rewrite: handle realloc(NULL, ...)

---
 c2rust-analyze/src/rewrite/expr/convert.rs | 31 ++++++++++---
 c2rust-analyze/src/rewrite/expr/mir_op.rs  | 53 ++++++++++++++--------
 2 files changed, 58 insertions(+), 26 deletions(-)

diff --git a/c2rust-analyze/src/rewrite/expr/convert.rs b/c2rust-analyze/src/rewrite/expr/convert.rs
index ff5fec82c..e2857e97d 100644
--- a/c2rust-analyze/src/rewrite/expr/convert.rs
+++ b/c2rust-analyze/src/rewrite/expr/convert.rs
@@ -345,6 +345,7 @@ impl<'tcx> ConvertVisitor<'tcx> {
                 elem_size,
                 src_single,
                 dest_single,
+                option,
             } => {
                 // `realloc(p, n)` -> `Box::new(...)`
                 assert!(matches!(hir_rw, Rewrite::Identity));
@@ -364,9 +365,14 @@ impl<'tcx> ConvertVisitor<'tcx> {
                 }
                 let expr = match (src_single, dest_single) {
                     (false, false) => {
+                        let src = if option {
+                            "src_ptr.unwrap_or(Box::new([]))"
+                        } else {
+                            "src_ptr"
+                        };
                         stmts.push(Rewrite::Let1(
                             "mut dest_ptr".into(),
-                            Box::new(Rewrite::Text("Vec::from(src_ptr)".into())),
+                            Box::new(format_rewrite!("Vec::from({src})")),
                         ));
                         stmts.push(format_rewrite!(
                             "dest_ptr.resize_with(dest_n, || {})",
@@ -375,8 +381,9 @@ impl<'tcx> ConvertVisitor<'tcx> {
                         Rewrite::Text("dest_ptr.into_boxed_slice()".into())
                     }
                     (false, true) => {
+                        let opt_flatten = if option { ".flatten()" } else { "" };
                         format_rewrite!(
-                            "src_ptr.into_iter().next().unwrap_or_else(|| {})",
+                            "src_ptr.into_iter(){opt_flatten}.next().unwrap_or_else(|| {})",
                             zeroize_expr
                         )
                     }
@@ -385,16 +392,28 @@ impl<'tcx> ConvertVisitor<'tcx> {
                             "mut dest_ptr".into(),
                             Box::new(Rewrite::Text("Vec::with_capacity(dest_n)".into())),
                         ));
-                        stmts.push(Rewrite::Text(
-                            "if dest_n >= 1 { dest_ptr.push(*src_ptr); }".into(),
-                        ));
+                        if option {
+                            stmts.push(Rewrite::Text(
+                                "if dest_n >= 1 { if let Some(src) = src_ptr { dest_ptr.push(*src); } }".into(),
+                            ));
+                        } else {
+                            stmts.push(Rewrite::Text(
+                                "if dest_n >= 1 { dest_ptr.push(*src_ptr); }".into(),
+                            ));
+                        }
                         stmts.push(format_rewrite!(
                             "dest_ptr.resize_with(dest_n, || {})",
                             zeroize_expr,
                         ));
                         Rewrite::Text("dest_ptr.into_boxed_slice()".into())
                     }
-                    (true, true) => Rewrite::Text("src_ptr".into()),
+                    (true, true) => {
+                        if option {
+                            format_rewrite!("src_ptr.unwrap_or_else(|| Box::new({}))", zeroize_expr)
+                        } else {
+                            Rewrite::Text("src_ptr".into())
+                        }
+                    }
                 };
                 Rewrite::Block(stmts, Some(Box::new(expr)))
             }
diff --git a/c2rust-analyze/src/rewrite/expr/mir_op.rs b/c2rust-analyze/src/rewrite/expr/mir_op.rs
index 9d9569edd..16d673790 100644
--- a/c2rust-analyze/src/rewrite/expr/mir_op.rs
+++ b/c2rust-analyze/src/rewrite/expr/mir_op.rs
@@ -111,6 +111,7 @@ pub enum RewriteKind {
         elem_size: u64,
         src_single: bool,
         dest_single: bool,
+        option: bool,
     },
     CallocSafe {
         zero_ty: ZeroizeType,
@@ -810,17 +811,25 @@ impl<'a, 'tcx> ExprRewriteVisitor<'a, 'tcx> {
                             };
 
                             // Cast input to either `Box<T>` or `Box<[T]>`, as in `free`.
+                            let mut option = false;
                             v.enter_call_arg(0, |v| {
-                                v.emit_cast_lty_adjust(src_lty, |desc| TypeDesc {
-                                    own: Ownership::Box,
-                                    qty: if src_single {
-                                        Quantity::Single
-                                    } else {
-                                        Quantity::Slice
-                                    },
-                                    dyn_owned: false,
-                                    option: desc.option,
-                                    pointee_ty: desc.pointee_ty,
+                                v.emit_cast_lty_adjust(src_lty, |desc| {
+                                    // `realloc(NULL, ...)` is explicitly allowed by the spec, so
+                                    // we can't force an unwrap here by returning `option: false`.
+                                    // Instead, we record the `option` flag as part of the rewrite
+                                    // so the nullable case can be handled appropriately.
+                                    option = desc.option;
+                                    TypeDesc {
+                                        own: Ownership::Box,
+                                        qty: if src_single {
+                                            Quantity::Single
+                                        } else {
+                                            Quantity::Slice
+                                        },
+                                        dyn_owned: false,
+                                        option: desc.option,
+                                        pointee_ty: desc.pointee_ty,
+                                    }
                                 });
                             });
 
@@ -829,21 +838,25 @@ impl<'a, 'tcx> ExprRewriteVisitor<'a, 'tcx> {
                                 elem_size,
                                 src_single,
                                 dest_single,
+                                option,
                             });
 
                             // Cast output from `Box<T>`/`Box<[T]>` to the target type, as in
                             // `malloc`.
                             v.emit_cast_adjust_lty(
-                                |desc| TypeDesc {
-                                    own: Ownership::Box,
-                                    qty: if dest_single {
-                                        Quantity::Single
-                                    } else {
-                                        Quantity::Slice
-                                    },
-                                    dyn_owned: false,
-                                    option: false,
-                                    pointee_ty: desc.pointee_ty,
+                                |desc| {
+                                    TypeDesc {
+                                        own: Ownership::Box,
+                                        qty: if dest_single {
+                                            Quantity::Single
+                                        } else {
+                                            Quantity::Slice
+                                        },
+                                        dyn_owned: false,
+                                        // We always return non-null from `realloc`.
+                                        option: false,
+                                        pointee_ty: desc.pointee_ty,
+                                    }
                                 },
                                 dest_lty,
                             );