[PATCH 4/7] ntdll: Make realloc implementation more generic.

[Rémi Bernon]

From: Rémi Bernon

Either resizing a block in place, or reallocating through RtlHeap calls.
---
 dlls/ntdll/heap.c | 74 +++++++++++++++--------------------------------
 1 file changed, 24 insertions(+), 50 deletions(-)

diff --git a/dlls/ntdll/heap.c b/dlls/ntdll/heap.c
index d2c5ee1ff3a..31650433aa5 100644
--- a/dlls/ntdll/heap.c
+++ b/dlls/ntdll/heap.c
@@ -869,31 +869,6 @@ static void free_large_block( struct heap *heap, struct block *block )
 }

-/***********************************************************************
- *           realloc_large_block
- */
-static struct block *realloc_large_block( struct heap *heap, DWORD flags, struct block *block, SIZE_T size )
-{
-    ARENA_LARGE *arena = CONTAINING_RECORD( block, ARENA_LARGE, block );
-    SIZE_T old_size = arena->data_size;
-    char *data = (char *)(block + 1);
-
-    if (flags & HEAP_REALLOC_IN_PLACE_ONLY) return NULL;
-    if (!(block = allocate_large_block( heap, flags, size )))
-    {
-        WARN( "Could not allocate block for %#Ix bytes\n", size );
-        return NULL;
-    }
-
-    valgrind_notify_alloc( block + 1, size, 0 );
-    memcpy( block + 1, data, old_size );
-    valgrind_notify_free( data );
-    free_large_block( heap, &arena->block );
-
-    return block;
-}
-
-
 /***********************************************************************
  *           find_large_block
  */
@@ -1622,15 +1597,13 @@ BOOLEAN WINAPI DECLSPEC_HOTPATCH RtlFreeHeap( HANDLE handle, ULONG flags, void *
     return !status;
 }

-
-static NTSTATUS heap_reallocate( struct heap *heap, ULONG flags, struct block *block, SIZE_T block_size,
-                                 SIZE_T size, SIZE_T *old_size, void **ret )
+static NTSTATUS heap_resize_block( struct heap *heap, ULONG flags, struct block *block, SIZE_T block_size,
+                                   SIZE_T size, SIZE_T *old_size, void **ret )
 {
     SUBHEAP *subheap = block_get_subheap( heap, block );
     SIZE_T old_block_size;
     struct entry *entry;
     struct block *next;
-    NTSTATUS status;

if (block_get_flags( block ) & BLOCK_FLAG_LARGE)
     {
@@ -1638,12 +1611,8 @@ static NTSTATUS heap_reallocate( struct heap *heap, ULONG flags, struct block *b
         old_block_size = large->block_size;
         *old_size = large->data_size;

-        if (old_block_size - sizeof(*block) < size)
-        {
-            if (!(block = realloc_large_block( heap, flags, block, size ))) return STATUS_NO_MEMORY;
-            *ret = block + 1;
-            return STATUS_SUCCESS;
-        }
+        if (block_size < HEAP_MIN_LARGE_BLOCK_SIZE / 4) return STATUS_NO_MEMORY;  /* shrinking large block to small block */
+        if (old_block_size < block_size) return STATUS_NO_MEMORY;

/* FIXME: we could remap zero-pages instead */
         valgrind_notify_resize( block + 1, *old_size, size );
@@ -1657,25 +1626,20 @@ static NTSTATUS heap_reallocate( struct heap *heap, ULONG flags, struct block *b
         return STATUS_SUCCESS;
     }

-    /* Check if we need to grow the block */
-
     old_block_size = block_get_size( block );
     *old_size = old_block_size - block_get_overhead( block );

+    if (block_size >= HEAP_MIN_LARGE_BLOCK_SIZE) return STATUS_NO_MEMORY;  /* growing small block to large block */
+
+    heap_lock( heap, flags );
+
     if (block_size > old_block_size)
     {
         if (!(next = next_block( subheap, block )) || !(block_get_flags( next ) & BLOCK_FLAG_FREE) ||
-            block_size >= HEAP_MIN_LARGE_BLOCK_SIZE ||
             block_size > old_block_size + block_get_size( next ) || !subheap_commit( heap, subheap, block, block_size ))
         {
-            if (flags & HEAP_REALLOC_IN_PLACE_ONLY) return STATUS_NO_MEMORY;
-            if ((status = heap_allocate( heap, flags & ~HEAP_ZERO_MEMORY, block_size, size, ret ))) return status;
-            valgrind_notify_alloc( *ret, size, 0 );
-            memcpy( *ret, block + 1, *old_size );
-            if (flags & HEAP_ZERO_MEMORY) memset( (char *)*ret + *old_size, 0, size - *old_size );
-            valgrind_notify_free( block + 1 );
-            free_used_block( heap, flags, block );
-            return STATUS_SUCCESS;
+            heap_unlock( heap, flags );
+            return STATUS_NO_MEMORY;
         }
     }

@@ -1694,6 +1658,8 @@ static NTSTATUS heap_reallocate( struct heap *heap, ULONG flags, struct block *b
     initialize_block( block, *old_size, size, flags );
     mark_block_tail( block, flags );

+    heap_unlock( heap, flags );
+
     *ret = block + 1;
     return STATUS_SUCCESS;
 }
@@ -1718,11 +1684,19 @@ void *WINAPI RtlReAllocateHeap( HANDLE handle, ULONG flags, void *ptr, SIZE_T si
         status = STATUS_NO_MEMORY;
     else if (!(block = unsafe_block_from_ptr( heap, heap_flags, ptr )))
         status = STATUS_INVALID_PARAMETER;
-    else
+    else if ((status = heap_resize_block( heap, heap_flags, block, block_size, size,
+                                          &old_size, &ret )))
     {
-        heap_lock( heap, heap_flags );
-        status = heap_reallocate( heap, heap_flags, block, block_size, size, &old_size, &ret );
-        heap_unlock( heap, heap_flags );
+        if (flags & HEAP_REALLOC_IN_PLACE_ONLY)
+            status = STATUS_NO_MEMORY;
+        else if (!(ret = RtlAllocateHeap( heap, flags, size )))
+            status = STATUS_NO_MEMORY;
+        else
+        {
+            memcpy( ret, ptr, min( size, old_size ) );
+            RtlFreeHeap( heap, flags, ptr );
+            status = STATUS_SUCCESS;
+        }
     }

TRACE( "handle %p, flags %#lx, ptr %p, size %#Ix, return %p, status %#lx.\n", handle, flags, ptr, size, ret, status );