base/allocator: Remove branches from the fast path in PartitionAlloc.

This significantly shortens the fast path in the allocator:

See the annotated assembly dump below for the fast path (release build,
x86_64):
- Before: 5 tests (and branches, shown with "***"), epilogue at offset
  0x8c
- After: 3 tests, epilogue at 0x46

BEFORE:
0000000000d385c0 <base::PartitionRoot<true>::AllocFromBucket(base::internal::PartitionBucket<true>*, int, unsigned long)>:
[...]  // Prologue
  d385c7:       48 83 ec 10             sub    $0x10,%rsp
  d385cb:       41 89 d6                mov    %edx,%r14d
  d385ce:       c6 45 ef 00             movb   $0x0,-0x11(%rbp)
  d385d2:       48 8b 16                mov    (%rsi),%rdx
  d385d5:       48 8b 1a                mov    (%rdx),%rbx
***  d385d8:       48 85 db                test   %rbx,%rbx
  d385db:       74 7b                   je     d38658 <base::PartitionRoot<true>::AllocFromBucket(base::internal::PartitionBucket<true>*, int, unsigned long)+0x98>
  d385dd:       48 8b 03                mov    (%rbx),%rax
  d385e0:       48 0f c8                bswap  %rax
  d385e3:       48 89 02                mov    %rax,(%rdx)
  d385e6:       66 83 42 18 01          addw   $0x1,0x18(%rdx)
***  d385eb:       48 85 db                test   %rbx,%rbx
  d385ee:       74 5c                   je     d3864c <base::PartitionRoot<true>::AllocFromBucket(base::internal::PartitionBucket<true>*, int, unsigned long)+0x8c>
  d385f0:       48 89 d9                mov    %rbx,%rcx
  d385f3:       48 81 e1 00 00 e0 ff    and    $0xffffffffffe00000,%rcx
  d385fa:       48 89 da                mov    %rbx,%rdx
  d385fd:       48 c1 ea 09             shr    $0x9,%rdx
  d38601:       81 e2 e0 0f 00 00       and    $0xfe0,%edx
  d38607:       48 8d 04 11             lea    (%rcx,%rdx,1),%rax
  d3860b:       48 05 00 10 00 00       add    $0x1000,%rax
  d38611:       0f b7 8c 11 1c 10 00    movzwl 0x101c(%rcx,%rdx,1),%ecx
  d38618:       00
  d38619:       48 c1 e1 05             shl    $0x5,%rcx
  d3861d:       48 89 c2                mov    %rax,%rdx
  d38620:       48 29 ca                sub    %rcx,%rdx
  d38623:       48 8b 52 10             mov    0x10(%rdx),%rdx
  d38627:       8b 72 18                mov    0x18(%rdx),%esi
***  d3862a:       48 81 fe 01 00 01 00    cmp    $0x10001,%rsi
  d38631:       73 48                   jae    d3867b <base::PartitionRoot<true>::AllocFromBucket(base::internal::PartitionBucket<true>*, int, unsigned long)+0xbb>
  d38633:       48 89 f2                mov    %rsi,%rdx
***  d38636:       41 f6 c6 02             test   $0x2,%r14b
  d3863a:       74 10                   je     d3864c <base::PartitionRoot<true>::AllocFromBucket(base::internal::PartitionBucket<true>*, int, unsigned long)+0x8c>
***  d3863c:       80 7d ef 00             cmpb   $0x0,-0x11(%rbp)
  d38640:       75 0a                   jne    d3864c <base::PartitionRoot<true>::AllocFromBucket(base::internal::PartitionBucket<true>*, int, unsigned long)+0x8c>
  d38642:       48 89 df                mov    %rbx,%rdi
  d38645:       31 f6                   xor    %esi,%esi
  d38647:       e8 54 e2 43 00          callq  11768a0 <memset@plt>
  d3864c:       48 89 d8                mov    %rbx,%rax
  [...]  // Epilogue
  d38657:       c3                      retq

AFTER:
0000000000d385c0 <base::PartitionRoot<true>::AllocFromBucket(base::internal::PartitionBucket<true>*, int, unsigned long)>:
[...] // Prologue
  d385ca:       49 89 cf                mov    %rcx,%r15
  d385cd:       41 89 d6                mov    %edx,%r14d
  d385d0:       c6 45 e7 00             movb   $0x0,-0x19(%rbp)
  d385d4:       48 8b 0e                mov    (%rsi),%rcx
  d385d7:       48 8b 19                mov    (%rcx),%rbx
***  d385da:       48 85 db                test   %rbx,%rbx
  d385dd:       74 35                   je     d38614 <base::PartitionRoot<true>::AllocFromBucket(base::internal::PartitionBucket<true>*, int, unsigned long)+0x54>
  d385df:       8b 56 18                mov    0x18(%rsi),%edx
  d385e2:       48 8b 03                mov    (%rbx),%rax
  d385e5:       48 0f c8                bswap  %rax
  d385e8:       48 89 01                mov    %rax,(%rcx)
  d385eb:       66 83 41 18 01          addw   $0x1,0x18(%rcx)
***  d385f0:       41 f6 c6 02             test   $0x2,%r14b
  d385f4:       74 10                   je     d38606 <base::PartitionRoot<true>::AllocFromBucket(base::internal::PartitionBucket<true>*, int, unsigned long)+0x46>
***  d385f6:       80 7d e7 00             cmpb   $0x0,-0x19(%rbp)
  d385fa:       75 0a                   jne    d38606 <base::PartitionRoot<true>::AllocFromBucket(base::internal::PartitionBucket<true>*, int, unsigned long)+0x46>
  d385fc:       48 89 df                mov    %rbx,%rdi
  d385ff:       31 f6                   xor    %esi,%esi
  d38601:       e8 9a e2 43 00          callq  11768a0 <memset@plt>
  d38606:       48 89 d8                mov    %rbx,%rax
[...]  // Epilogue
  d38613:       c3                      retq

Bug: 998048
Change-Id: Ib9698e635b459b90c86a2dabadb3e4536f1d2219
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2307333Reviewed-by: Kentaro Hara <haraken@chromium.org>
Commit-Queue: Benoit L <lizeb@chromium.org>
Cr-Commit-Position: refs/heads/master@{#790285}

base/allocator/partition_allocator/partition_alloc.h

@@ -536,6 +536,8 @@ ALWAYS_INLINE void* PartitionRoot<thread_safe>::AllocFromBucket(Bucket* bucket,
   // Check that this page is neither full nor freed.
   PA_DCHECK(page);
   PA_DCHECK(page->num_allocated_slots >= 0);
+  size_t new_slot_size = bucket->slot_size;
+
   void* ret = page->freelist_head;
   if (LIKELY(ret)) {
     // If these DCHECKs fire, you probably corrupted memory. TODO(palmer): See
@@ -550,24 +552,23 @@ ALWAYS_INLINE void* PartitionRoot<thread_safe>::AllocFromBucket(Bucket* bucket,
             page->freelist_head->next);
     page->freelist_head = new_head;
     page->num_allocated_slots++;
+
+    PA_DCHECK(page->bucket == bucket);
   } else {
     ret = bucket->SlowPathAlloc(this, flags, size, &is_already_zeroed);
     // TODO(palmer): See if we can afford to make this a CHECK.
     PA_DCHECK(!ret || IsValidPage(Page::FromPointer(ret)));
-  }
 
-  if (!ret) {
-    return nullptr;
-  }
+    if (UNLIKELY(!ret))
+      return nullptr;
 
-  page = Page::FromPointer(ret);
-  // TODO(ajwong): Can |page->bucket| ever not be |bucket|? If not, can this
-  // just be bucket->slot_size?
-  size_t new_slot_size = page->bucket->slot_size;
-  size_t raw_size = page->get_raw_size();
-  if (raw_size) {
-    PA_DCHECK(raw_size == size);
-    new_slot_size = raw_size;
+    page = Page::FromPointer(ret);
+    if (UNLIKELY(page->get_raw_size())) {
+      PA_DCHECK(page->get_raw_size() == size);
+      new_slot_size = size;
+    } else {
+      new_slot_size = page->bucket->slot_size;
+    }
   }
   // Layout inside the slot: |[tag]|cookie|object|[empty]|cookie|
   //                                      <--a--->