tdbBtree.c 23.3 KB
Newer Older
H
Hongze Cheng 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13
/*
 * Copyright (c) 2019 TAOS Data, Inc. <jhtao@taosdata.com>
 *
 * This program is free software: you can use, redistribute, and/or modify
 * it under the terms of the GNU Affero General Public License, version 3
 * or later ("AGPL"), as published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
H
refact  
Hongze Cheng 已提交
14 15
 */

H
Hongze Cheng 已提交
16 17
#include "tdbInt.h"

H
Hongze Cheng 已提交
18 19 20 21 22
#define TDB_BTREE_ROOT 0x1
#define TDB_BTREE_LEAF 0x2

#define TDB_BTREE_PAGE_IS_ROOT(flags) TDB_FLAG_HAS(flags, TDB_BTREE_ROOT)
#define TDB_BTREE_PAGE_IS_LEAF(flags) TDB_FLAG_HAS(flags, TDB_BTREE_LEAF)
H
Hongze Cheng 已提交
23 24 25
#define TDB_BTREE_ASSERT_FLAG(flags)                                                 \
  ASSERT(TDB_FLAG_IS(flags, TDB_BTREE_ROOT) || TDB_FLAG_IS(flags, TDB_BTREE_LEAF) || \
         TDB_FLAG_IS(flags, TDB_BTREE_ROOT | TDB_BTREE_LEAF) || TDB_FLAG_IS(flags, 0))
H
Hongze Cheng 已提交
26

H
Hongze Cheng 已提交
27
struct SBTree {
H
more  
Hongze Cheng 已提交
28 29 30
  SPgno          root;
  int            keyLen;
  int            valLen;
H
Hongze Cheng 已提交
31
  SPager        *pPager;
H
more  
Hongze Cheng 已提交
32
  FKeyComparator kcmpr;
H
more  
Hongze Cheng 已提交
33
  u8             fanout;
H
Hongze Cheng 已提交
34
  int            pageSize;
H
more  
Hongze Cheng 已提交
35 36 37 38
  int            maxLocal;
  int            minLocal;
  int            maxLeaf;
  int            minLeaf;
H
Hongze Cheng 已提交
39
  u8            *pTmp;
H
Hongze Cheng 已提交
40 41
};

H
Hongze Cheng 已提交
42 43
#define TDB_BTREE_PAGE_COMMON_HDR u8 flags;

H
Hongze Cheng 已提交
44 45
#define TDB_BTREE_PAGE_GET_FLAGS(PAGE)        (PAGE)->pData[0]
#define TDB_BTREE_PAGE_SET_FLAGS(PAGE, flags) ((PAGE)->pData[0] = (flags))
H
Hongze Cheng 已提交
46 47 48 49 50

typedef struct __attribute__((__packed__)) {
  TDB_BTREE_PAGE_COMMON_HDR
} SLeafHdr;

H
Hongze Cheng 已提交
51
typedef struct __attribute__((__packed__)) {
H
Hongze Cheng 已提交
52 53 54
  TDB_BTREE_PAGE_COMMON_HDR;
  SPgno pgno;  // right-most child
} SIntHdr;
H
more  
Hongze Cheng 已提交
55

H
Hongze Cheng 已提交
56
typedef struct {
H
Hongze Cheng 已提交
57
  u8      flags;
H
Hongze Cheng 已提交
58
  SBTree *pBt;
H
Hongze Cheng 已提交
59
} SBtreeInitPageArg;
H
Hongze Cheng 已提交
60

H
Hongze Cheng 已提交
61 62 63 64 65 66 67 68 69
typedef struct {
  int   kLen;
  u8   *pKey;
  int   vLen;
  u8   *pVal;
  SPgno pgno;
  u8   *pTmpSpace;
} SCellDecoder;

H
Hongze Cheng 已提交
70
static int tdbBtCursorMoveTo(SBtCursor *pCur, const void *pKey, int kLen, int *pCRst);
H
Hongze Cheng 已提交
71
static int tdbDefaultKeyCmprFn(const void *pKey1, int keyLen1, const void *pKey2, int keyLen2);
H
Hongze Cheng 已提交
72
static int tdbBtreeOpenImpl(SBTree *pBt);
H
Hongze Cheng 已提交
73 74
static int tdbBtreeZeroPage(SPage *pPage, void *arg);
static int tdbBtreeInitPage(SPage *pPage, void *arg);
H
Hongze Cheng 已提交
75 76
static int tdbBtreeEncodeCell(SPage *pPage, const void *pKey, int kLen, const void *pVal, int vLen, SCell *pCell,
                              int *szCell);
H
Hongze Cheng 已提交
77
static int tdbBtreeDecodeCell(SPage *pPage, const SCell *pCell, SCellDecoder *pDecoder);
H
Hongze Cheng 已提交
78
static int tdbBtreeBalance(SBtCursor *pCur);
H
Hongze Cheng 已提交
79
static int tdbBtreeCellSize(const SPage *pPage, SCell *pCell);
H
Hongze Cheng 已提交
80

H
Hongze Cheng 已提交
81
int tdbBtreeOpen(int keyLen, int valLen, SPager *pPager, FKeyComparator kcmpr, SBTree **ppBt) {
H
Hongze Cheng 已提交
82
  SBTree *pBt;
H
Hongze Cheng 已提交
83
  int     ret;
H
more  
Hongze Cheng 已提交
84

H
Hongze Cheng 已提交
85
  *ppBt = NULL;
H
Hongze Cheng 已提交
86 87 88 89 90 91 92 93 94 95

  pBt = (SBTree *)calloc(1, sizeof(*pBt));
  if (pBt == NULL) {
    return -1;
  }

  // pBt->keyLen
  pBt->keyLen = keyLen;
  // pBt->valLen
  pBt->valLen = valLen;
H
refact  
Hongze Cheng 已提交
96 97
  // pBt->pPager
  pBt->pPager = pPager;
H
Hongze Cheng 已提交
98 99
  // pBt->kcmpr
  pBt->kcmpr = kcmpr ? kcmpr : tdbDefaultKeyCmprFn;
H
more  
Hongze Cheng 已提交
100 101 102 103 104 105 106
  // pBt->fanout
  if (keyLen == TDB_VARIANT_LEN) {
    pBt->fanout = TDB_DEFAULT_FANOUT;
  } else {
    ASSERT(0);
    // TODO: pBt->fanout = 0;
  }
H
Hongze Cheng 已提交
107 108 109
  // pBt->pageSize
  pBt->pageSize = tdbPagerGetPageSize(pPager);
  // pBt->maxLocal
H
Hongze Cheng 已提交
110
  pBt->maxLocal = (pBt->pageSize - 14) / pBt->fanout;
H
Hongze Cheng 已提交
111
  // pBt->minLocal: Should not be allowed smaller than 15, which is [nPayload][nKey][nData]
H
Hongze Cheng 已提交
112
  pBt->minLocal = (pBt->pageSize - 14) / pBt->fanout / 2;
H
Hongze Cheng 已提交
113
  // pBt->maxLeaf
H
Hongze Cheng 已提交
114
  pBt->maxLeaf = pBt->pageSize - 14;
H
Hongze Cheng 已提交
115 116
  // pBt->minLeaf
  pBt->minLeaf = pBt->minLocal;
H
Hongze Cheng 已提交
117

H
Hongze Cheng 已提交
118 119 120 121 122 123 124
  // TODO: pBt->root
  ret = tdbBtreeOpenImpl(pBt);
  if (ret < 0) {
    free(pBt);
    return -1;
  }

H
Hongze Cheng 已提交
125
  *ppBt = pBt;
H
Hongze Cheng 已提交
126 127 128 129 130 131 132 133
  return 0;
}

int tdbBtreeClose(SBTree *pBt) {
  // TODO
  return 0;
}

H
more  
Hongze Cheng 已提交
134
int tdbBtreeCursor(SBtCursor *pCur, SBTree *pBt) {
H
more  
Hongze Cheng 已提交
135 136
  pCur->pBt = pBt;
  pCur->iPage = -1;
H
more  
Hongze Cheng 已提交
137
  pCur->pPage = NULL;
H
Hongze Cheng 已提交
138
  pCur->idx = -1;
H
more  
Hongze Cheng 已提交
139

H
more  
Hongze Cheng 已提交
140 141 142
  return 0;
}

H
Hongze Cheng 已提交
143
int tdbBtCursorInsert(SBtCursor *pCur, const void *pKey, int kLen, const void *pVal, int vLen) {
H
Hongze Cheng 已提交
144
  int     ret;
H
Hongze Cheng 已提交
145
  int     idx;
H
Hongze Cheng 已提交
146
  SPager *pPager;
H
Hongze Cheng 已提交
147 148 149
  SCell  *pCell;
  int     szCell;
  int     cret;
H
Hongze Cheng 已提交
150
  SBTree *pBt;
H
Hongze Cheng 已提交
151

H
Hongze Cheng 已提交
152
  ret = tdbBtCursorMoveTo(pCur, pKey, kLen, &cret);
H
Hongze Cheng 已提交
153 154 155 156 157
  if (ret < 0) {
    // TODO: handle error
    return -1;
  }

H
Hongze Cheng 已提交
158
  if (pCur->idx == -1) {
H
Hongze Cheng 已提交
159
    ASSERT(TDB_PAGE_TOTAL_CELLS(pCur->pPage) == 0);
H
Hongze Cheng 已提交
160 161 162
    idx = 0;
  } else {
    if (cret > 0) {
H
Hongze Cheng 已提交
163
      idx = pCur->idx + 1;
H
Hongze Cheng 已提交
164
    } else if (cret < 0) {
H
Hongze Cheng 已提交
165
      idx = pCur->idx;
H
Hongze Cheng 已提交
166 167 168 169 170 171
    } else {
      /* TODO */
      ASSERT(0);
    }
  }

H
Hongze Cheng 已提交
172 173 174 175 176 177 178 179 180 181 182
  // TODO: refact code here
  pBt = pCur->pBt;
  if (!pBt->pTmp) {
    pBt->pTmp = (u8 *)malloc(pBt->pageSize);
    if (pBt->pTmp == NULL) {
      return -1;
    }
  }

  pCell = pBt->pTmp;

H
Hongze Cheng 已提交
183
  // Encode the cell
H
Hongze Cheng 已提交
184 185 186 187
  ret = tdbBtreeEncodeCell(pCur->pPage, pKey, kLen, pVal, vLen, pCell, &szCell);
  if (ret < 0) {
    return -1;
  }
H
Hongze Cheng 已提交
188 189 190 191 192

  // Insert the cell to the index
  ret = tdbPageInsertCell(pCur->pPage, idx, pCell, szCell);
  if (ret < 0) {
    return -1;
H
Hongze Cheng 已提交
193 194
  }

H
Hongze Cheng 已提交
195 196 197 198 199 200 201 202
  // If page is overflow, balance the tree
  if (pCur->pPage->nOverflow > 0) {
    ret = tdbBtreeBalance(pCur);
    if (ret < 0) {
      return -1;
    }
  }

H
more  
Hongze Cheng 已提交
203 204 205
  return 0;
}

H
Hongze Cheng 已提交
206 207 208 209 210
static int tdbBtCursorMoveToChild(SBtCursor *pCur, SPgno pgno) {
  // TODO
  return 0;
}

H
Hongze Cheng 已提交
211
static int tdbBtCursorMoveTo(SBtCursor *pCur, const void *pKey, int kLen, int *pCRst) {
H
Hongze Cheng 已提交
212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231
  int     ret;
  SBTree *pBt;
  SPager *pPager;

  pBt = pCur->pBt;
  pPager = pBt->pPager;

  if (pCur->iPage < 0) {
    ASSERT(pCur->iPage == -1);
    ASSERT(pCur->idx == -1);

    // Move from the root
    ret = tdbPagerFetchPage(pPager, pBt->root, &(pCur->pPage), tdbBtreeInitPage, pBt);
    if (ret < 0) {
      ASSERT(0);
      return -1;
    }

    pCur->iPage = 0;

H
Hongze Cheng 已提交
232
    if (TDB_PAGE_TOTAL_CELLS(pCur->pPage) == 0) {
H
Hongze Cheng 已提交
233
      // Current page is empty
H
Hongze Cheng 已提交
234
      // ASSERT(TDB_FLAG_IS(TDB_PAGE_FLAGS(pCur->pPage), TDB_BTREE_ROOT | TDB_BTREE_LEAF));
H
Hongze Cheng 已提交
235 236 237
      return 0;
    }

H
Hongze Cheng 已提交
238
    for (;;) {
H
Hongze Cheng 已提交
239
      int          lidx, ridx, midx, c, nCells;
H
Hongze Cheng 已提交
240 241 242 243 244
      SCell       *pCell;
      SPage       *pPage;
      SCellDecoder cd = {0};

      pPage = pCur->pPage;
H
Hongze Cheng 已提交
245
      nCells = TDB_PAGE_TOTAL_CELLS(pPage);
H
Hongze Cheng 已提交
246
      lidx = 0;
H
Hongze Cheng 已提交
247 248 249 250
      ridx = nCells - 1;

      ASSERT(nCells > 0);

H
Hongze Cheng 已提交
251
      for (;;) {
H
Hongze Cheng 已提交
252 253 254 255
        if (lidx > ridx) break;

        midx = (lidx + ridx) >> 1;

H
Hongze Cheng 已提交
256
        pCell = tdbPageGetCell(pPage, midx);
H
Hongze Cheng 已提交
257 258 259 260 261 262
        ret = tdbBtreeDecodeCell(pPage, pCell, &cd);
        if (ret < 0) {
          // TODO: handle error
          ASSERT(0);
          return -1;
        }
H
Hongze Cheng 已提交
263

H
Hongze Cheng 已提交
264 265 266
        // Compare the key values
        c = pBt->kcmpr(pKey, kLen, cd.pKey, cd.kLen);
        if (c < 0) {
H
Hongze Cheng 已提交
267 268
          /* input-key < cell-key */
          ridx = midx - 1;
H
Hongze Cheng 已提交
269
        } else if (c > 0) {
H
Hongze Cheng 已提交
270 271
          /* input-key > cell-key */
          lidx = midx + 1;
H
Hongze Cheng 已提交
272
        } else {
H
Hongze Cheng 已提交
273 274 275 276 277
          /* input-key == cell-key */
          break;
        }
      }

H
Hongze Cheng 已提交
278
      // Move downward or break
H
Hongze Cheng 已提交
279 280
      u8 flags = TDB_BTREE_PAGE_GET_FLAGS(pPage);
      u8 leaf = TDB_BTREE_PAGE_IS_LEAF(flags);
H
Hongze Cheng 已提交
281 282
      if (leaf) {
        pCur->idx = midx;
H
Hongze Cheng 已提交
283
        *pCRst = c;
H
Hongze Cheng 已提交
284 285 286 287 288 289 290 291 292
        break;
      } else {
        if (c <= 0) {
          pCur->idx = midx;
          tdbBtCursorMoveToChild(pCur, cd.pgno);
        } else {
          if (midx == nCells - 1) {
            /* Move to right-most child */
            pCur->idx = midx + 1;
H
Hongze Cheng 已提交
293
            // tdbBtCursorMoveToChild(pCur, ((SBtPageHdr *)(pPage->pAmHdr))->rChild);
H
Hongze Cheng 已提交
294 295 296
          } else {
            // TODO: reset cd as uninitialized
            pCur->idx = midx + 1;
H
Hongze Cheng 已提交
297
            pCell = tdbPageGetCell(pPage, midx + 1);
H
Hongze Cheng 已提交
298 299 300
            tdbBtreeDecodeCell(pPage, pCell, &cd);
            tdbBtCursorMoveToChild(pCur, cd.pgno);
          }
H
Hongze Cheng 已提交
301 302
        }
      }
H
Hongze Cheng 已提交
303 304 305 306
    }

  } else {
    // TODO: Move the cursor from a some position instead of a clear state
H
Hongze Cheng 已提交
307
    ASSERT(0);
H
Hongze Cheng 已提交
308
  }
H
more  
Hongze Cheng 已提交
309

H
Hongze Cheng 已提交
310 311
  return 0;
}
H
more  
Hongze Cheng 已提交
312

H
Hongze Cheng 已提交
313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330
static int tdbDefaultKeyCmprFn(const void *pKey1, int keyLen1, const void *pKey2, int keyLen2) {
  int mlen;
  int cret;

  ASSERT(keyLen1 > 0 && keyLen2 > 0 && pKey1 != NULL && pKey2 != NULL);

  mlen = keyLen1 < keyLen2 ? keyLen1 : keyLen2;
  cret = memcmp(pKey1, pKey2, mlen);
  if (cret == 0) {
    if (keyLen1 < keyLen2) {
      cret = -1;
    } else if (keyLen1 > keyLen2) {
      cret = 1;
    } else {
      cret = 0;
    }
  }
  return cret;
H
Hongze Cheng 已提交
331 332 333 334 335 336 337 338 339 340
}

static int tdbBtreeOpenImpl(SBTree *pBt) {
  // Try to get the root page of the an existing btree

  SPgno  pgno;
  SPage *pPage;
  int    ret;

  {
H
Hongze Cheng 已提交
341
    // 1. TODO: Search the main DB to check if the DB exists
H
Hongze Cheng 已提交
342 343 344 345 346 347 348 349 350
    pgno = 0;
  }

  if (pgno != 0) {
    pBt->root = pgno;
    return 0;
  }

  // Try to create a new database
H
Hongze Cheng 已提交
351
  SBtreeInitPageArg zArg = {.flags = TDB_BTREE_ROOT | TDB_BTREE_LEAF, .pBt = pBt};
H
Hongze Cheng 已提交
352 353 354 355 356 357
  ret = tdbPagerNewPage(pBt->pPager, &pgno, &pPage, tdbBtreeZeroPage, &zArg);
  if (ret < 0) {
    return -1;
  }

  // TODO: Unref the page
H
Hongze Cheng 已提交
358

H
Hongze Cheng 已提交
359 360
  ASSERT(pgno != 0);
  pBt->root = pgno;
H
Hongze Cheng 已提交
361

H
Hongze Cheng 已提交
362 363 364
  return 0;
}

H
Hongze Cheng 已提交
365
static int tdbBtreeInitPage(SPage *pPage, void *arg) {
H
Hongze Cheng 已提交
366
  SBTree *pBt;
H
Hongze Cheng 已提交
367
  u8      flags;
H
Hongze Cheng 已提交
368
  u8      isLeaf;
H
Hongze Cheng 已提交
369

H
Hongze Cheng 已提交
370 371
  pBt = (SBTree *)arg;
  flags = TDB_BTREE_PAGE_GET_FLAGS(pPage);
H
Hongze Cheng 已提交
372
  isLeaf = TDB_BTREE_PAGE_IS_LEAF(flags);
H
Hongze Cheng 已提交
373

H
Hongze Cheng 已提交
374
  ASSERT(flags == TDB_BTREE_PAGE_GET_FLAGS(pPage));
H
Hongze Cheng 已提交
375

H
Hongze Cheng 已提交
376
  tdbPageInit(pPage, isLeaf ? sizeof(SLeafHdr) : sizeof(SIntHdr), tdbBtreeCellSize);
H
Hongze Cheng 已提交
377

H
Hongze Cheng 已提交
378 379
  TDB_BTREE_ASSERT_FLAG(flags);

H
Hongze Cheng 已提交
380
  if (isLeaf) {
H
Hongze Cheng 已提交
381 382 383 384 385 386 387 388 389 390
    pPage->kLen = pBt->keyLen;
    pPage->vLen = pBt->valLen;
    pPage->maxLocal = pBt->maxLeaf;
    pPage->minLocal = pBt->minLeaf;
  } else {
    pPage->kLen = pBt->keyLen;
    pPage->vLen = sizeof(SPgno);
    pPage->maxLocal = pBt->maxLocal;
    pPage->minLocal = pBt->minLocal;
  }
H
Hongze Cheng 已提交
391 392 393 394

  return 0;
}

H
Hongze Cheng 已提交
395
static int tdbBtreeZeroPage(SPage *pPage, void *arg) {
H
Hongze Cheng 已提交
396
  u8      flags;
H
Hongze Cheng 已提交
397
  SBTree *pBt;
H
Hongze Cheng 已提交
398
  u8      isLeaf;
H
Hongze Cheng 已提交
399

H
Hongze Cheng 已提交
400 401
  flags = ((SBtreeInitPageArg *)arg)->flags;
  pBt = ((SBtreeInitPageArg *)arg)->pBt;
H
Hongze Cheng 已提交
402
  isLeaf = TDB_BTREE_PAGE_IS_LEAF(flags);
H
Hongze Cheng 已提交
403

H
Hongze Cheng 已提交
404
  tdbPageZero(pPage, isLeaf ? sizeof(SLeafHdr) : sizeof(SIntHdr), tdbBtreeCellSize);
H
Hongze Cheng 已提交
405 406

  if (isLeaf) {
H
Hongze Cheng 已提交
407
    SLeafHdr *pLeafHdr = (SLeafHdr *)(pPage->pData);
H
Hongze Cheng 已提交
408 409
    pLeafHdr->flags = flags;

H
Hongze Cheng 已提交
410 411 412 413 414
    pPage->kLen = pBt->keyLen;
    pPage->vLen = pBt->valLen;
    pPage->maxLocal = pBt->maxLeaf;
    pPage->minLocal = pBt->minLeaf;
  } else {
H
Hongze Cheng 已提交
415
    SIntHdr *pIntHdr = (SIntHdr *)(pPage->pData);
H
Hongze Cheng 已提交
416 417 418
    pIntHdr->flags = flags;
    pIntHdr->pgno = 0;

H
Hongze Cheng 已提交
419 420 421 422 423
    pPage->kLen = pBt->keyLen;
    pPage->vLen = sizeof(SPgno);
    pPage->maxLocal = pBt->maxLocal;
    pPage->minLocal = pBt->minLocal;
  }
H
Hongze Cheng 已提交
424

H
Hongze Cheng 已提交
425
  return 0;
H
Hongze Cheng 已提交
426 427 428
}

#ifndef TDB_BTREE_BALANCE
H
Hongze Cheng 已提交
429 430 431 432
typedef struct {
  SBTree *pBt;
  SPage  *pParent;
  int     idx;
H
Hongze Cheng 已提交
433
  i8      nOld;
H
Hongze Cheng 已提交
434 435 436 437 438
  SPage  *pOldPages[3];
  i8      nNewPages;
  SPage  *pNewPages[5];
} SBtreeBalanceHelper;

H
Hongze Cheng 已提交
439 440 441 442 443
static int tdbBtreeBalanceDeeper(SBTree *pBt, SPage *pRoot, SPage **ppChild) {
  SPager           *pPager;
  SPage            *pChild;
  SPgno             pgnoChild;
  int               ret;
H
Hongze Cheng 已提交
444
  SIntHdr          *pIntHdr;
H
Hongze Cheng 已提交
445
  SBtreeInitPageArg zArg;
H
Hongze Cheng 已提交
446 447 448 449 450 451 452 453 454 455 456 457

  pPager = pRoot->pPager;

  // Allocate a new child page
  zArg.flags = TDB_BTREE_LEAF;
  zArg.pBt = pBt;
  ret = tdbPagerNewPage(pPager, &pgnoChild, &pChild, tdbBtreeZeroPage, &zArg);
  if (ret < 0) {
    return -1;
  }

  // Copy the root page content to the child page
H
Hongze Cheng 已提交
458
  tdbPageCopy(pRoot, pChild);
H
Hongze Cheng 已提交
459 460 461 462 463 464 465 466 467

  // Reinitialize the root page
  zArg.flags = TDB_BTREE_ROOT;
  zArg.pBt = pBt;
  ret = tdbBtreeZeroPage(pRoot, &zArg);
  if (ret < 0) {
    return -1;
  }

H
Hongze Cheng 已提交
468
  pIntHdr = (SIntHdr *)(pRoot->pData);
H
Hongze Cheng 已提交
469
  pIntHdr->pgno = pgnoChild;
H
Hongze Cheng 已提交
470

H
Hongze Cheng 已提交
471 472 473 474
  *ppChild = pChild;
  return 0;
}

H
Hongze Cheng 已提交
475
static int tdbBtreeBalanceNonRoot(SBTree *pBt, SPage *pParent, int idx) {
H
Hongze Cheng 已提交
476 477 478 479
  int ret;

  int    nOlds;
  SPage *pOlds[3];
H
Hongze Cheng 已提交
480 481
  SCell *pDivCell[2] = {0};
  int    szDivCell[2];
H
Hongze Cheng 已提交
482
  int    sIdx;
H
Hongze Cheng 已提交
483
  u8     childLeaf;
H
Hongze Cheng 已提交
484

H
Hongze Cheng 已提交
485
  {  // Find 3 child pages at most to do balance
H
Hongze Cheng 已提交
486 487 488
    int    nCells = TDB_PAGE_TOTAL_CELLS(pParent);
    SCell *pCell;

H
Hongze Cheng 已提交
489 490 491 492 493 494 495 496 497 498 499 500 501 502
    if (nCells <= 2) {
      sIdx = 0;
      nOlds = nCells + 1;
    } else {
      // has more than three child pages
      if (idx == 0) {
        sIdx = 0;
      } else if (idx == nCells) {
        sIdx = idx - 2;
      } else {
        sIdx = idx - 1;
      }
      nOlds = 3;
    }
H
Hongze Cheng 已提交
503 504
    for (int i = 0; i < nOlds; i++) {
      ASSERT(sIdx + i <= nCells);
H
Hongze Cheng 已提交
505

H
Hongze Cheng 已提交
506
      SPgno pgno;
H
Hongze Cheng 已提交
507
      if (sIdx + i == nCells) {
H
Hongze Cheng 已提交
508 509 510
        ASSERT(!TDB_BTREE_PAGE_IS_LEAF(TDB_BTREE_PAGE_GET_FLAGS(pParent)));
        pgno = ((SIntHdr *)(pParent->pData))->pgno;
      } else {
H
Hongze Cheng 已提交
511
        pCell = tdbPageGetCell(pParent, sIdx + i);
H
Hongze Cheng 已提交
512 513 514 515 516 517 518 519 520
        pgno = *(SPgno *)pCell;
      }

      ret = tdbPagerFetchPage(pBt->pPager, pgno, pOlds + i, tdbBtreeInitPage, pBt);
      if (ret < 0) {
        ASSERT(0);
        return -1;
      }
    }
H
Hongze Cheng 已提交
521 522 523 524 525 526 527 528 529 530 531
    // copy the parent key out if child pages are not leaf page
    childLeaf = TDB_BTREE_PAGE_IS_LEAF(TDB_BTREE_PAGE_GET_FLAGS(pOlds[0]));
    if (childLeaf) {
      for (int i = 0; i < nOlds - 1; i++) {
        pCell = tdbPageGetCell(pParent, sIdx + i);

        szDivCell[i] = tdbBtreeCellSize(pParent, pCell);
        pDivCell[i] = malloc(szDivCell[i]);
        memcpy(pDivCell, pCell, szDivCell[i]);
      }
    }
H
Hongze Cheng 已提交
532
    // drop the cells on parent page
H
Hongze Cheng 已提交
533 534 535 536 537 538 539 540
    for (int i = 0; i < nOlds; i++) {
      nCells = TDB_PAGE_TOTAL_CELLS(pParent);
      if (sIdx < nCells) {
        tdbPageDropCell(pParent, sIdx);
      } else {
        ((SIntHdr *)pParent->pData)->pgno = 0;
      }
    }
H
Hongze Cheng 已提交
541 542
  }

H
Hongze Cheng 已提交
543
  int nNews = 0;
H
Hongze Cheng 已提交
544 545 546 547 548 549
  struct {
    int cnt;
    int size;
    int oPage;
    int oIdx;
  } infoNews[5] = {0};
H
Hongze Cheng 已提交
550 551 552

  {  // Get how many new pages are needed and the new distribution

H
Hongze Cheng 已提交
553 554 555
    // first loop to find minimum number of pages needed
    for (int oPage = 0; oPage < nOlds; oPage++) {
      SPage *pPage = pOlds[oPage];
H
Hongze Cheng 已提交
556 557
      SCell *pCell;
      int    cellBytes;
H
Hongze Cheng 已提交
558
      int    oIdx;
H
Hongze Cheng 已提交
559

H
Hongze Cheng 已提交
560
      for (oIdx = 0; oIdx < TDB_PAGE_TOTAL_CELLS(pPage); oIdx++) {
H
Hongze Cheng 已提交
561
        pCell = tdbPageGetCell(pPage, oIdx);
H
Hongze Cheng 已提交
562 563
        cellBytes = TDB_BYTES_CELL_TAKEN(pPage, pCell);

H
Hongze Cheng 已提交
564 565
        if (infoNews[nNews].size + cellBytes > TDB_PAGE_USABLE_SIZE(pPage)) {
          // page is full, use a new page
H
Hongze Cheng 已提交
566
          nNews++;
H
Hongze Cheng 已提交
567
          // for a child leaf case, this cell is used as the new divider cell to parent
H
Hongze Cheng 已提交
568 569 570 571 572 573 574 575 576 577 578 579
          if (childLeaf) continue;
        }
        infoNews[nNews].cnt++;
        infoNews[nNews].size += cellBytes;
        infoNews[nNews].oPage = oPage;
        infoNews[nNews].oIdx = oIdx;
      }

      // For child leaf pages
      if (childLeaf && oPage < nOlds - 1) {
        if (infoNews[nNews].size + szDivCell[oPage] + TDB_PAGE_OFFSET_SIZE(pPage) > TDB_PAGE_USABLE_SIZE(pPage)) {
          nNews++;
H
Hongze Cheng 已提交
580
        }
H
Hongze Cheng 已提交
581 582 583 584
        infoNews[nNews].cnt++;
        infoNews[nNews].size += cellBytes;
        infoNews[nNews].oPage = oPage;
        infoNews[nNews].oIdx = oIdx;
H
Hongze Cheng 已提交
585 586 587
      }
    }

H
Hongze Cheng 已提交
588 589 590 591
    nNews = nNews + 1;

    // back loop to make the distribution even
    for (int iNew = nNews - 1; iNew > 0; iNew--) {
H
Hongze Cheng 已提交
592
      SCell *pCell;
H
Hongze Cheng 已提交
593
      SPage *pPage;
H
Hongze Cheng 已提交
594
      int    nCells;
H
Hongze Cheng 已提交
595
      int    cellBytes;
H
Hongze Cheng 已提交
596

H
Hongze Cheng 已提交
597 598 599
      pPage = pOlds[infoNews[iNew - 1].oPage];
      nCells = TDB_PAGE_TOTAL_CELLS(pPage);

H
Hongze Cheng 已提交
600
      for (;;) {
H
Hongze Cheng 已提交
601
        pCell = tdbPageGetCell(pPage, infoNews[iNew - 1].oIdx);
H
Hongze Cheng 已提交
602 603
        cellBytes = TDB_BYTES_CELL_TAKEN(pPage, pCell);

H
Hongze Cheng 已提交
604 605 606 607 608 609
        infoNews[iNew].cnt++;
        infoNews[iNew].size += cellBytes;
        infoNews[iNew - 1].cnt--;
        infoNews[iNew - 1].size -= cellBytes;
        if ((infoNews[iNew - 1].oIdx--) == 0) {
          infoNews[iNew - 1].oPage--;
H
Hongze Cheng 已提交
610 611 612
          pPage = pOlds[infoNews[iNew - 1].oPage];
          nCells = TDB_PAGE_TOTAL_CELLS(pPage);
          infoNews[iNew - 1].oIdx = TDB_PAGE_TOTAL_CELLS(pPage);
H
Hongze Cheng 已提交
613 614
        }

H
Hongze Cheng 已提交
615 616
        if (infoNews[iNew].size > infoNews[iNew - 1].size) {
          break;
H
Hongze Cheng 已提交
617 618
        }
      }
H
Hongze Cheng 已提交
619
    }
H
Hongze Cheng 已提交
620 621
  }

H
Hongze Cheng 已提交
622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642
  SPage *pNews[5] = {0};
  {  // Allocate new pages, reuse the old page when possible

    SPgno             pgno;
    SBtreeInitPageArg iarg;
    u8                flags;

    flags = TDB_BTREE_PAGE_GET_FLAGS(pOlds[0]);

    for (int iNew = 0; iNew < nNews; iNew++) {
      if (iNew < nOlds) {
        pNews[iNew] = pOlds[iNew];
      } else {
        iarg.pBt = pBt;
        iarg.flags = flags;
        ret = tdbPagerNewPage(pBt->pPager, &pgno, pNews + iNew, tdbBtreeZeroPage, &iarg);
        if (ret < 0) {
          ASSERT(0);
        }
      }
    }
H
Hongze Cheng 已提交
643 644

    // TODO: sort the page according to the page number
H
Hongze Cheng 已提交
645 646
  }

H
Hongze Cheng 已提交
647
  {  // Do the real cell distribution
H
Hongze Cheng 已提交
648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692

    SPage            *pTPage[2];
    int               tPage, tIdx, iOld;
    SCell            *pCell;
    int               szCell;
    SBtreeInitPageArg iarg = {.flags = TDB_BTREE_PAGE_GET_FLAGS(pOlds[0]), .pBt = pBt};

    for (int i = 0; i < 2; i++) {
      ret = tdbPageCreate(pOlds[0]->pageSize, &pTPage[i], NULL, NULL);
      if (ret < 0) {
        ASSERT(0);
      }
    }

    tPage = 0;
    tIdx = 0;
    iOld = 0;
    tdbBtreeZeroPage(pTPage[tPage], &iarg);
    tdbPageCopy(pOlds[iOld++], pTPage[tPage]);

    for (int iNew = 0; iNew < nNews; iNew++) {
      // fill the iNew page
      tdbBtreeZeroPage(pNews[iNew], &iarg);

      for (int iCell = 0; iCell < infoNews[iNew].cnt; iCell++) {
        while (tIdx >= TDB_PAGE_TOTAL_CELLS(pTPage[tPage])) {
          tPage = (tPage + 1) % 2;
          tIdx = 0;

          tdbBtreeZeroPage(pTPage[tPage], &iarg);
          tdbPageCopy(pOlds[iOld++], pTPage[tPage]);
        }

        pCell = tdbPageGetCell(pTPage[tPage], tIdx);
        szCell = tdbBtreeCellSize(pTPage[tPage], pCell);

        tdbPageInsertCell(pNews[iNew], iCell, pCell, szCell);

        tIdx++;
      }
    }

    for (int i = 0; i < 2; i++) {
      tdbPageDestroy(pTPage[i], NULL, NULL);
    }
H
Hongze Cheng 已提交
693
  }
H
Hongze Cheng 已提交
694

H
Hongze Cheng 已提交
695
  {  // Insert records in parent page
H
Hongze Cheng 已提交
696
    int          cIdx = sIdx;
H
Hongze Cheng 已提交
697 698 699 700 701
    int          szCell;
    SCell        pCell[128];  // TODO
    SCellDecoder cd;

    for (int iNew = 0; iNew < nNews; iNew++) {
H
Hongze Cheng 已提交
702 703 704 705 706 707 708 709 710 711
      if (iNew == nNews - 1) {
        // The last new page

        SIntHdr *pIntHdr = (SIntHdr *)pParent->pData;
        if (pIntHdr->pgno == 0) {
          pIntHdr->pgno = TDB_PAGE_PGNO(pNews[iNew]);
          break;
        }
      }

H
Hongze Cheng 已提交
712 713 714
      tdbBtreeDecodeCell(pNews[iNew], tdbPageGetCell(pNews[iNew], TDB_PAGE_TOTAL_CELLS(pNews[iNew]) - 1), &cd);

      tdbBtreeEncodeCell(pParent, cd.pKey, cd.kLen, (void *)&TDB_PAGE_PGNO(pNews[iNew]), sizeof(SPgno), pCell, &szCell);
H
Hongze Cheng 已提交
715
      // TODO: the cell here may be used by pParent as an overflow cell
H
Hongze Cheng 已提交
716 717 718 719
      tdbPageInsertCell(pParent, cIdx, pCell, szCell);
    }
  }

H
Hongze Cheng 已提交
720 721 722 723 724
  return 0;
}

static int tdbBtreeBalance(SBtCursor *pCur) {
  int    iPage;
H
Hongze Cheng 已提交
725
  SPage *pParent;
H
Hongze Cheng 已提交
726
  SPage *pPage;
H
Hongze Cheng 已提交
727
  int    ret;
H
Hongze Cheng 已提交
728
  u8     flags;
H
Hongze Cheng 已提交
729 730
  u8     leaf;
  u8     root;
H
Hongze Cheng 已提交
731 732 733 734

  // Main loop to balance the BTree
  for (;;) {
    iPage = pCur->iPage;
H
Hongze Cheng 已提交
735
    pPage = pCur->pPage;
H
Hongze Cheng 已提交
736
    flags = TDB_BTREE_PAGE_GET_FLAGS(pPage);
H
Hongze Cheng 已提交
737 738
    leaf = TDB_BTREE_PAGE_IS_LEAF(flags);
    root = TDB_BTREE_PAGE_IS_ROOT(flags);
H
Hongze Cheng 已提交
739

H
Hongze Cheng 已提交
740 741
    // when the page is not overflow and not too empty, the balance work
    // is finished. Just break out the balance loop.
H
Hongze Cheng 已提交
742
    if (pPage->nOverflow == 0 /* TODO: && pPage->nFree <= */) {
H
Hongze Cheng 已提交
743 744 745 746
      break;
    }

    if (iPage == 0) {
H
Hongze Cheng 已提交
747 748 749
      // For the root page, only balance when the page is overfull,
      // ignore the case of empty
      if (pPage->nOverflow == 0) break;
H
Hongze Cheng 已提交
750

H
Hongze Cheng 已提交
751
      ret = tdbBtreeBalanceDeeper(pCur->pBt, pPage, &(pCur->pgStack[1]));
H
Hongze Cheng 已提交
752 753 754 755 756 757 758 759 760
      if (ret < 0) {
        return -1;
      }

      pCur->idx = 0;
      pCur->idxStack[0] = 0;
      pCur->pgStack[0] = pCur->pPage;
      pCur->iPage = 1;
      pCur->pPage = pCur->pgStack[1];
H
Hongze Cheng 已提交
761
    } else {
H
Hongze Cheng 已提交
762
      // Generalized balance step
H
Hongze Cheng 已提交
763
      pParent = pCur->pgStack[iPage - 1];
H
Hongze Cheng 已提交
764

H
Hongze Cheng 已提交
765
      ret = tdbBtreeBalanceNonRoot(pCur->pBt, pParent, pCur->idxStack[pCur->iPage - 1]);
H
Hongze Cheng 已提交
766 767 768 769 770 771
      if (ret < 0) {
        return -1;
      }

      pCur->iPage--;
      pCur->pPage = pCur->pgStack[pCur->iPage];
H
Hongze Cheng 已提交
772 773 774 775 776
    }
  }

  return 0;
}
H
Hongze Cheng 已提交
777 778
#endif

H
Hongze Cheng 已提交
779
#ifndef TDB_BTREE_CELL  // =========================================================
H
Hongze Cheng 已提交
780
static int tdbBtreeEncodePayload(SPage *pPage, u8 *pPayload, const void *pKey, int kLen, const void *pVal, int vLen,
H
Hongze Cheng 已提交
781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809
                                 int *szPayload) {
  int nPayload;

  ASSERT(pKey != NULL);

  if (pVal == NULL) {
    vLen = 0;
  }

  nPayload = kLen + vLen;
  if (nPayload <= pPage->maxLocal) {
    // General case without overflow
    memcpy(pPayload, pKey, kLen);
    if (pVal) {
      memcpy(pPayload + kLen, pVal, vLen);
    }

    *szPayload = nPayload;
    return 0;
  }

  {
    // TODO: handle overflow case
    ASSERT(0);
  }

  return 0;
}

H
Hongze Cheng 已提交
810 811
static int tdbBtreeEncodeCell(SPage *pPage, const void *pKey, int kLen, const void *pVal, int vLen, SCell *pCell,
                              int *szCell) {
H
Hongze Cheng 已提交
812
  u8  flags;
H
Hongze Cheng 已提交
813 814 815 816
  u8  leaf;
  int nHeader;
  int nPayload;
  int ret;
H
Hongze Cheng 已提交
817 818 819 820

  ASSERT(pPage->kLen == TDB_VARIANT_LEN || pPage->kLen == kLen);
  ASSERT(pPage->vLen == TDB_VARIANT_LEN || pPage->vLen == vLen);

H
Hongze Cheng 已提交
821 822
  nPayload = 0;
  nHeader = 0;
H
Hongze Cheng 已提交
823
  flags = TDB_BTREE_PAGE_GET_FLAGS(pPage);
H
Hongze Cheng 已提交
824
  leaf = TDB_BTREE_PAGE_IS_LEAF(flags);
H
Hongze Cheng 已提交
825

H
Hongze Cheng 已提交
826
  // 1. Encode Header part
H
Hongze Cheng 已提交
827 828 829 830 831 832 833 834
  /* Encode SPgno if interior page */
  if (!leaf) {
    ASSERT(pPage->vLen == sizeof(SPgno));

    ((SPgno *)(pCell + nHeader))[0] = ((SPgno *)pVal)[0];
    nHeader = nHeader + sizeof(SPgno);
  }

H
Hongze Cheng 已提交
835
  /* Encode kLen if need */
H
Hongze Cheng 已提交
836
  if (pPage->kLen == TDB_VARIANT_LEN) {
H
Hongze Cheng 已提交
837
    nHeader += tdbPutVarInt(pCell + nHeader, kLen);
H
Hongze Cheng 已提交
838 839
  }

H
Hongze Cheng 已提交
840
  /* Encode vLen if need */
H
Hongze Cheng 已提交
841
  if (pPage->vLen == TDB_VARIANT_LEN) {
H
Hongze Cheng 已提交
842
    nHeader += tdbPutVarInt(pCell + nHeader, vLen);
H
Hongze Cheng 已提交
843 844
  }

H
Hongze Cheng 已提交
845 846 847 848 849 850 851 852 853
  // 2. Encode payload part
  if (leaf) {
    ret = tdbBtreeEncodePayload(pPage, pCell + nHeader, pKey, kLen, pVal, vLen, &nPayload);
  } else {
    ret = tdbBtreeEncodePayload(pPage, pCell + nHeader, pKey, kLen, NULL, 0, &nPayload);
  }
  if (ret < 0) {
    // TODO: handle error
    return -1;
H
Hongze Cheng 已提交
854 855
  }

H
Hongze Cheng 已提交
856
  *szCell = nHeader + nPayload;
H
Hongze Cheng 已提交
857 858 859
  return 0;
}

H
Hongze Cheng 已提交
860 861 862 863
static int tdbBtreeDecodePayload(SPage *pPage, const u8 *pPayload, SCellDecoder *pDecoder) {
  int nPayload;

  ASSERT(pDecoder->pKey == NULL);
H
Hongze Cheng 已提交
864

H
Hongze Cheng 已提交
865 866
  if (pDecoder->pVal) {
    nPayload = pDecoder->kLen + pDecoder->vLen;
H
Hongze Cheng 已提交
867
  } else {
H
Hongze Cheng 已提交
868
    nPayload = pDecoder->kLen;
H
Hongze Cheng 已提交
869 870
  }

H
Hongze Cheng 已提交
871 872 873
  if (nPayload <= pPage->maxLocal) {
    // General case without overflow
    pDecoder->pKey = (void *)pPayload;
H
Hongze Cheng 已提交
874 875
    if (!pDecoder->pVal) {
      pDecoder->pVal = (void *)(pPayload + pDecoder->kLen);
H
Hongze Cheng 已提交
876
    }
H
Hongze Cheng 已提交
877
  } else {
H
Hongze Cheng 已提交
878 879
    // TODO: handle overflow case
    ASSERT(0);
H
Hongze Cheng 已提交
880 881
  }

H
Hongze Cheng 已提交
882 883 884 885
  return 0;
}

static int tdbBtreeDecodeCell(SPage *pPage, const SCell *pCell, SCellDecoder *pDecoder) {
H
Hongze Cheng 已提交
886
  u8  flags;
H
Hongze Cheng 已提交
887 888 889 890 891
  u8  leaf;
  int nHeader;
  int ret;

  nHeader = 0;
H
Hongze Cheng 已提交
892
  flags = TDB_BTREE_PAGE_GET_FLAGS(pPage);
H
Hongze Cheng 已提交
893 894
  leaf = TDB_BTREE_PAGE_IS_LEAF(flags);

H
Hongze Cheng 已提交
895 896 897 898 899 900 901
  // Clear the state of decoder
  pDecoder->kLen = -1;
  pDecoder->pKey = NULL;
  pDecoder->vLen = -1;
  pDecoder->pVal = NULL;
  pDecoder->pgno = 0;

H
Hongze Cheng 已提交
902
  // 1. Decode header part
H
Hongze Cheng 已提交
903 904 905 906 907 908 909 910
  if (!leaf) {
    ASSERT(pPage->vLen == sizeof(SPgno));

    pDecoder->pgno = ((SPgno *)(pCell + nHeader))[0];
    pDecoder->pVal = (u8 *)(&(pDecoder->pgno));
    nHeader = nHeader + sizeof(SPgno);
  }

H
Hongze Cheng 已提交
911 912 913 914 915 916 917 918
  if (pPage->kLen == TDB_VARIANT_LEN) {
    nHeader += tdbGetVarInt(pCell + nHeader, &(pDecoder->kLen));
  } else {
    pDecoder->kLen = pPage->kLen;
  }

  if (pPage->vLen == TDB_VARIANT_LEN) {
    nHeader += tdbGetVarInt(pCell + nHeader, &(pDecoder->vLen));
H
Hongze Cheng 已提交
919
  } else {
H
Hongze Cheng 已提交
920 921
    pDecoder->vLen = pPage->vLen;
  }
H
Hongze Cheng 已提交
922

H
Hongze Cheng 已提交
923 924 925 926
  // 2. Decode payload part
  ret = tdbBtreeDecodePayload(pPage, pCell + nHeader, pDecoder);
  if (ret < 0) {
    return -1;
H
Hongze Cheng 已提交
927 928 929 930 931
  }

  return 0;
}

H
Hongze Cheng 已提交
932
static int tdbBtreeCellSize(const SPage *pPage, SCell *pCell) {
H
Hongze Cheng 已提交
933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965
  u8  flags;
  u8  isLeaf;
  int szCell;
  int kLen = 0, vLen = 0;

  flags = TDB_BTREE_PAGE_GET_FLAGS(pPage);
  isLeaf = TDB_BTREE_PAGE_IS_LEAF(flags);
  szCell = 0;

  if (!isLeaf) {
    szCell += sizeof(SPgno);
  }

  if (pPage->kLen == TDB_VARIANT_LEN) {
    szCell += tdbGetVarInt(pCell + szCell, &kLen);
  } else {
    kLen = pPage->kLen;
  }

  if (isLeaf) {
    if (pPage->vLen == TDB_VARIANT_LEN) {
      szCell += tdbGetVarInt(pCell + szCell, &vLen);
    } else {
      vLen = pPage->vLen;
    }
  }

  szCell = szCell + kLen + vLen;

  return szCell;
}

#endif