tdbBtree.c 19.9 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 480
  int ret;

  int    nOlds;
  SPage *pOlds[3];

H
Hongze Cheng 已提交
481
  {  // Find 3 child pages at most to do balance
H
Hongze Cheng 已提交
482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499
    int nCells = TDB_PAGE_TOTAL_CELLS(pParent);
    int sIdx;
    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;
    }
    for (int i = 0; i < nOlds; i++, sIdx++) {
      ASSERT(sIdx <= nCells);
H
Hongze Cheng 已提交
500

H
Hongze Cheng 已提交
501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519
      SPgno pgno;
      if (sIdx == nCells) {
        ASSERT(!TDB_BTREE_PAGE_IS_LEAF(TDB_BTREE_PAGE_GET_FLAGS(pParent)));
        pgno = ((SIntHdr *)(pParent->pData))->pgno;
      } else {
        SCell *pCell;

        pCell = tdbPageGetCell(pParent, sIdx);
        pgno = *(SPgno *)pCell;
      }

      ret = tdbPagerFetchPage(pBt->pPager, pgno, pOlds + i, tdbBtreeInitPage, pBt);
      if (ret < 0) {
        ASSERT(0);
        return -1;
      }
    }
  }

H
Hongze Cheng 已提交
520
  int nNews = 0;
H
Hongze Cheng 已提交
521 522 523 524 525 526
  struct {
    int cnt;
    int size;
    int oPage;
    int oIdx;
  } infoNews[5] = {0};
H
Hongze Cheng 已提交
527 528 529

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

H
Hongze Cheng 已提交
530 531 532
    // first loop to find minimum number of pages needed
    for (int oPage = 0; oPage < nOlds; oPage++) {
      SPage *pPage = pOlds[oPage];
H
Hongze Cheng 已提交
533 534
      SCell *pCell;
      int    cellBytes;
H
Hongze Cheng 已提交
535

H
Hongze Cheng 已提交
536 537
      for (int oIdx = 0; oIdx < TDB_PAGE_TOTAL_CELLS(pPage); oIdx++) {
        pCell = tdbPageGetCell(pPage, oIdx);
H
Hongze Cheng 已提交
538 539
        cellBytes = TDB_BYTES_CELL_TAKEN(pPage, pCell);

H
Hongze Cheng 已提交
540 541
        if (infoNews[nNews].size + cellBytes > TDB_PAGE_USABLE_SIZE(pPage)) {
          // page is full, use a new page
H
Hongze Cheng 已提交
542 543
          nNews++;
        }
H
Hongze Cheng 已提交
544 545 546 547
        infoNews[nNews].cnt++;
        infoNews[nNews].size += cellBytes;
        infoNews[nNews].oPage = oPage;
        infoNews[nNews].oIdx = oIdx;
H
Hongze Cheng 已提交
548 549 550
      }
    }

H
Hongze Cheng 已提交
551 552 553 554
    nNews = nNews + 1;

    // back loop to make the distribution even
    for (int iNew = nNews - 1; iNew > 0; iNew--) {
H
Hongze Cheng 已提交
555
      SCell *pCell;
H
Hongze Cheng 已提交
556 557
      SPage *pPage;
      int    cellBytes;
H
Hongze Cheng 已提交
558 559

      for (;;) {
H
Hongze Cheng 已提交
560 561
        pPage = pOlds[infoNews[iNew - 1].oPage];
        pCell = tdbPageGetCell(pPage, infoNews[iNew - 1].oIdx);
H
Hongze Cheng 已提交
562 563
        cellBytes = TDB_BYTES_CELL_TAKEN(pPage, pCell);

H
Hongze Cheng 已提交
564 565 566 567 568 569 570
        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--;
          infoNews[iNew - 1].oIdx = TDB_PAGE_TOTAL_CELLS(pOlds[infoNews[iNew - 1].oPage]);
H
Hongze Cheng 已提交
571 572
        }

H
Hongze Cheng 已提交
573 574
        if (infoNews[iNew].size > infoNews[iNew - 1].size) {
          break;
H
Hongze Cheng 已提交
575 576
        }
      }
H
Hongze Cheng 已提交
577
    }
H
Hongze Cheng 已提交
578 579
  }

H
Hongze Cheng 已提交
580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600
  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 已提交
601 602 603
  }

  {
H
Hongze Cheng 已提交
604
      // Do the actual redistribute
H
Hongze Cheng 已提交
605
  }
H
Hongze Cheng 已提交
606

H
Hongze Cheng 已提交
607 608 609 610 611
  return 0;
}

static int tdbBtreeBalance(SBtCursor *pCur) {
  int    iPage;
H
Hongze Cheng 已提交
612
  SPage *pParent;
H
Hongze Cheng 已提交
613
  SPage *pPage;
H
Hongze Cheng 已提交
614
  int    ret;
H
Hongze Cheng 已提交
615
  u8     flags;
H
Hongze Cheng 已提交
616 617
  u8     leaf;
  u8     root;
H
Hongze Cheng 已提交
618 619 620 621

  // Main loop to balance the BTree
  for (;;) {
    iPage = pCur->iPage;
H
Hongze Cheng 已提交
622
    pPage = pCur->pPage;
H
Hongze Cheng 已提交
623
    flags = TDB_BTREE_PAGE_GET_FLAGS(pPage);
H
Hongze Cheng 已提交
624 625
    leaf = TDB_BTREE_PAGE_IS_LEAF(flags);
    root = TDB_BTREE_PAGE_IS_ROOT(flags);
H
Hongze Cheng 已提交
626

H
Hongze Cheng 已提交
627 628
    // when the page is not overflow and not too empty, the balance work
    // is finished. Just break out the balance loop.
H
Hongze Cheng 已提交
629
    if (pPage->nOverflow == 0 /* TODO: && pPage->nFree <= */) {
H
Hongze Cheng 已提交
630 631 632 633
      break;
    }

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

H
Hongze Cheng 已提交
638
      ret = tdbBtreeBalanceDeeper(pCur->pBt, pPage, &(pCur->pgStack[1]));
H
Hongze Cheng 已提交
639 640 641 642 643 644 645 646 647
      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 已提交
648
    } else {
H
Hongze Cheng 已提交
649
      // Generalized balance step
H
Hongze Cheng 已提交
650
      pParent = pCur->pgStack[iPage - 1];
H
Hongze Cheng 已提交
651

H
Hongze Cheng 已提交
652
      ret = tdbBtreeBalanceNonRoot(pCur->pBt, pParent, pCur->idxStack[pCur->iPage - 1]);
H
Hongze Cheng 已提交
653 654 655 656 657 658
      if (ret < 0) {
        return -1;
      }

      pCur->iPage--;
      pCur->pPage = pCur->pgStack[pCur->iPage];
H
Hongze Cheng 已提交
659 660 661 662 663
    }
  }

  return 0;
}
H
Hongze Cheng 已提交
664 665
#endif

H
Hongze Cheng 已提交
666
#ifndef TDB_BTREE_CELL  // =========================================================
H
Hongze Cheng 已提交
667
static int tdbBtreeEncodePayload(SPage *pPage, u8 *pPayload, const void *pKey, int kLen, const void *pVal, int vLen,
H
Hongze Cheng 已提交
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 693 694 695 696
                                 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 已提交
697 698
static int tdbBtreeEncodeCell(SPage *pPage, const void *pKey, int kLen, const void *pVal, int vLen, SCell *pCell,
                              int *szCell) {
H
Hongze Cheng 已提交
699
  u8  flags;
H
Hongze Cheng 已提交
700 701 702 703
  u8  leaf;
  int nHeader;
  int nPayload;
  int ret;
H
Hongze Cheng 已提交
704 705 706 707

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

H
Hongze Cheng 已提交
708 709
  nPayload = 0;
  nHeader = 0;
H
Hongze Cheng 已提交
710
  flags = TDB_BTREE_PAGE_GET_FLAGS(pPage);
H
Hongze Cheng 已提交
711
  leaf = TDB_BTREE_PAGE_IS_LEAF(flags);
H
Hongze Cheng 已提交
712

H
Hongze Cheng 已提交
713
  // 1. Encode Header part
H
Hongze Cheng 已提交
714 715 716 717 718 719 720 721
  /* 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 已提交
722
  /* Encode kLen if need */
H
Hongze Cheng 已提交
723
  if (pPage->kLen == TDB_VARIANT_LEN) {
H
Hongze Cheng 已提交
724
    nHeader += tdbPutVarInt(pCell + nHeader, kLen);
H
Hongze Cheng 已提交
725 726
  }

H
Hongze Cheng 已提交
727
  /* Encode vLen if need */
H
Hongze Cheng 已提交
728
  if (pPage->vLen == TDB_VARIANT_LEN) {
H
Hongze Cheng 已提交
729
    nHeader += tdbPutVarInt(pCell + nHeader, vLen);
H
Hongze Cheng 已提交
730 731
  }

H
Hongze Cheng 已提交
732 733 734 735 736 737 738 739 740
  // 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 已提交
741 742
  }

H
Hongze Cheng 已提交
743
  *szCell = nHeader + nPayload;
H
Hongze Cheng 已提交
744 745 746
  return 0;
}

H
Hongze Cheng 已提交
747 748 749 750
static int tdbBtreeDecodePayload(SPage *pPage, const u8 *pPayload, SCellDecoder *pDecoder) {
  int nPayload;

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

H
Hongze Cheng 已提交
752 753
  if (pDecoder->pVal) {
    nPayload = pDecoder->kLen + pDecoder->vLen;
H
Hongze Cheng 已提交
754
  } else {
H
Hongze Cheng 已提交
755
    nPayload = pDecoder->kLen;
H
Hongze Cheng 已提交
756 757
  }

H
Hongze Cheng 已提交
758 759 760
  if (nPayload <= pPage->maxLocal) {
    // General case without overflow
    pDecoder->pKey = (void *)pPayload;
H
Hongze Cheng 已提交
761 762
    if (!pDecoder->pVal) {
      pDecoder->pVal = (void *)(pPayload + pDecoder->kLen);
H
Hongze Cheng 已提交
763
    }
H
Hongze Cheng 已提交
764
  } else {
H
Hongze Cheng 已提交
765 766
    // TODO: handle overflow case
    ASSERT(0);
H
Hongze Cheng 已提交
767 768
  }

H
Hongze Cheng 已提交
769 770 771 772
  return 0;
}

static int tdbBtreeDecodeCell(SPage *pPage, const SCell *pCell, SCellDecoder *pDecoder) {
H
Hongze Cheng 已提交
773
  u8  flags;
H
Hongze Cheng 已提交
774 775 776 777 778
  u8  leaf;
  int nHeader;
  int ret;

  nHeader = 0;
H
Hongze Cheng 已提交
779
  flags = TDB_BTREE_PAGE_GET_FLAGS(pPage);
H
Hongze Cheng 已提交
780 781
  leaf = TDB_BTREE_PAGE_IS_LEAF(flags);

H
Hongze Cheng 已提交
782 783 784 785 786 787 788
  // Clear the state of decoder
  pDecoder->kLen = -1;
  pDecoder->pKey = NULL;
  pDecoder->vLen = -1;
  pDecoder->pVal = NULL;
  pDecoder->pgno = 0;

H
Hongze Cheng 已提交
789
  // 1. Decode header part
H
Hongze Cheng 已提交
790 791 792 793 794 795 796 797
  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 已提交
798 799 800 801 802 803 804 805
  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 已提交
806
  } else {
H
Hongze Cheng 已提交
807 808
    pDecoder->vLen = pPage->vLen;
  }
H
Hongze Cheng 已提交
809

H
Hongze Cheng 已提交
810 811 812 813
  // 2. Decode payload part
  ret = tdbBtreeDecodePayload(pPage, pCell + nHeader, pDecoder);
  if (ret < 0) {
    return -1;
H
Hongze Cheng 已提交
814 815 816 817 818
  }

  return 0;
}

H
Hongze Cheng 已提交
819
static int tdbBtreeCellSize(const SPage *pPage, SCell *pCell) {
H
Hongze Cheng 已提交
820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852
  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