From 3a9679b952ce5c2428a7f6c5e0b4038210502c0d Mon Sep 17 00:00:00 2001
From: Brandon Liu <bdon@bdon.org>
Date: Thu, 17 Feb 2022 17:16:44 +0800
Subject: [PATCH] handle leaf levels other than 7

---
 js/pmtiles.test.ts |  30 ++++++++++++
 js/pmtiles.ts      | 112 +++++++++++++++++++++++++--------------------
 2 files changed, 92 insertions(+), 50 deletions(-)

diff --git a/js/pmtiles.test.ts b/js/pmtiles.test.ts
index f2b54dc..38634c5 100644
--- a/js/pmtiles.test.ts
+++ b/js/pmtiles.test.ts
@@ -4,6 +4,7 @@ import {
   getUint24,
   getUint48,
   queryLeafdir,
+  queryLeafLevel,
   queryTile,
   parseEntry,
   Entry,
@@ -79,6 +80,35 @@ test("get leafdir", (assertion) => {
   assertion.ok(queryTile(view, 14, 16383, 16383) === null);
 });
 
+test("derive the leaf level", (assertion) => {
+  let data = createDirectory([
+    {
+      z: 6,
+      x: 3,
+      y: 3,
+      offset: 0,
+      length: 0,
+      is_dir: true,
+    },
+  ]);
+  let view = new DataView(data);
+  let level = queryLeafLevel(view);
+  assertion.ok(level === 6);
+  data = createDirectory([
+    {
+      z: 6,
+      x: 3,
+      y: 3,
+      offset: 0,
+      length: 0,
+      is_dir: false,
+    },
+  ]);
+  view = new DataView(data);
+  level = queryLeafLevel(view);
+  assertion.ok(level === null);
+});
+
 test("convert spec v1 directory to spec v2 directory", (assertion) => {
   let data = createDirectory([
     {
diff --git a/js/pmtiles.ts b/js/pmtiles.ts
index aa91ec8..35868c9 100644
--- a/js/pmtiles.ts
+++ b/js/pmtiles.ts
@@ -16,6 +16,39 @@ export const getUint48 = (view: DataView, pos: number) => {
   return shift(view.getUint32(pos + 2, true), 16) + view.getUint16(pos, true);
 };
 
+interface Zxy {
+  z: number;
+  x: number;
+  y: number;
+}
+
+interface Header {
+  version: number;
+  json_size: number;
+  root_entries: number;
+}
+
+interface Root {
+  header: Header;
+  buffer: ArrayBuffer;
+  dir: DataView;
+  // etag: string | null;
+}
+
+export interface Entry {
+  z: number;
+  x: number;
+  y: number;
+  offset: number;
+  length: number;
+  is_dir: boolean;
+}
+
+interface CachedLeaf {
+  lastUsed: number;
+  buffer: Promise<ArrayBuffer>;
+}
+
 const compare = (
   tz: number,
   tx: number,
@@ -88,6 +121,14 @@ const queryView = (
   return null;
 };
 
+export const queryLeafLevel = (view: DataView): number | null => {
+  if (view.byteLength < 17) return null;
+  let numEntries = view.byteLength / 17;
+  let entry = parseEntry(view, numEntries - 1);
+  if (entry.is_dir) return entry.z;
+  return null;
+};
+
 const entrySort = (a: Entry, b: Entry): number => {
   if (a.is_dir && !b.is_dir) {
     return 1;
@@ -159,47 +200,17 @@ export const createDirectory = (entries: Entry[]): ArrayBuffer => {
   return buffer;
 };
 
-interface Zxy {
-  z: number;
-  x: number;
-  y: number;
-}
-
-// TODO: handle different leaf levels
-export const deriveLeaf = (tile: Zxy): Zxy => {
-  let z7_tile_diff = tile.z - 7;
-  let z7_x = Math.trunc(tile.x / (1 << z7_tile_diff));
-  let z7_y = Math.trunc(tile.y / (1 << z7_tile_diff));
-  return { z: 7, x: z7_x, y: z7_y };
+export const deriveLeaf = (root: Root, tile: Zxy): Zxy | null => {
+  const leaf_level = queryLeafLevel(root.dir);
+  if (leaf_level) {
+    let level_diff = tile.z - leaf_level;
+    let leaf_x = Math.trunc(tile.x / (1 << level_diff));
+    let leaf_y = Math.trunc(tile.y / (1 << level_diff));
+    return { z: leaf_level, x: leaf_x, y: leaf_y };
+  }
+  return null;
 };
 
-interface Header {
-  version: number;
-  json_size: number;
-  root_entries: number;
-}
-
-interface Root {
-  header: Header;
-  buffer: ArrayBuffer;
-  dir: DataView;
-  // etag: string | null;
-}
-
-export interface Entry {
-  z: number;
-  x: number;
-  y: number;
-  offset: number;
-  length: number;
-  is_dir: boolean;
-}
-
-interface CachedLeaf {
-  lastUsed: number;
-  buffer: Promise<ArrayBuffer>;
-}
-
 export const parseHeader = (dataview: DataView): Header => {
   var magic = dataview.getUint16(0, true);
   if (magic !== 19792) {
@@ -329,17 +340,18 @@ export class PMTiles {
     let entry = queryTile(root.dir, z, x, y);
     if (entry) return entry;
 
-    let leafcoords = deriveLeaf({ z: z, x: x, y: y });
-    let leafdir_entry = queryLeafdir(
-      root.dir,
-      leafcoords.z,
-      leafcoords.x,
-      leafcoords.y
-    );
-
-    if (leafdir_entry) {
-      let leafdir = await this.getLeafdir(root.header.version, leafdir_entry);
-      return queryTile(new DataView(leafdir), z, x, y);
+    let leafcoords = deriveLeaf(root, { z: z, x: x, y: y });
+    if (leafcoords) {
+      let leafdir_entry = queryLeafdir(
+        root.dir,
+        leafcoords.z,
+        leafcoords.x,
+        leafcoords.y
+      );
+      if (leafdir_entry) {
+        let leafdir = await this.getLeafdir(root.header.version, leafdir_entry);
+        return queryTile(new DataView(leafdir), z, x, y);
+      }
     }
     return null;
   }