From 97bbc22a1d9d5b1b515e7f731a5084c575c0ac03 Mon Sep 17 00:00:00 2001
From: Pavel Safonov <p.safonov@postgrespro.ru>
Date: Thu, 22 Dec 2022 11:22:02 +0300
Subject: [PATCH] postgres: add page arg in pg_btree, change args names in
 pg_heap

---
 format/format.go                                   | 10 +++++++---
 .../postgres/common/pg_btree/postgres/pg_btree.go  |  8 +++++---
 format/postgres/common/pg_heap/postgres/pg_heap.go |  2 +-
 format/postgres/pg_btree.go                        | 14 ++++++++++----
 format/postgres/pg_heap.go                         |  6 +++---
 5 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/format/format.go b/format/format.go
index 3daadf254..a260c493b 100644
--- a/format/format.go
+++ b/format/format.go
@@ -332,7 +332,11 @@ type PostgresIn struct {
 }
 
 type PostgresHeapIn struct {
-	Flavour       string `doc:"PostgreSQL flavour: postgres14, pgproee14.., postgres10"`
-	PageNumber    int    `doc:"First page number in file, default is 0"`
-	SegmentNumber int    `doc:"Segment file number (16790.1 is 1), default is 0"`
+	Flavour string `doc:"PostgreSQL flavour: postgres14, pgproee14.., postgres10"`
+	Page    int    `doc:"First page number in file, default is 0"`
+	Segment int    `doc:"Segment file number (16790.1 is 1), default is 0"`
+}
+
+type PostgresBTreeIn struct {
+	Page int `doc:"First page number in file, default is 0"`
 }
diff --git a/format/postgres/common/pg_btree/postgres/pg_btree.go b/format/postgres/common/pg_btree/postgres/pg_btree.go
index 8ec52e8d0..e349ffb62 100644
--- a/format/postgres/common/pg_btree/postgres/pg_btree.go
+++ b/format/postgres/common/pg_btree/postgres/pg_btree.go
@@ -1,6 +1,7 @@
 package postgres
 
 import (
+	"github.com/wader/fq/format"
 	"github.com/wader/fq/format/postgres/common"
 	"github.com/wader/fq/format/postgres/common/pg_heap/postgres"
 	"github.com/wader/fq/pkg/decode"
@@ -66,8 +67,9 @@ const (
 // IndexTupleData *IndexTuple;
 /* total size (bytes):    8 */
 
-func DecodePgBTree(d *decode.D) any {
+func DecodePgBTree(d *decode.D, args format.PostgresBTreeIn) any {
 	btree := &BTree{
+		Args:     args,
 		PageSize: common.PageSize,
 	}
 	decodeBTreePages(btree, d)
@@ -75,13 +77,13 @@ func DecodePgBTree(d *decode.D) any {
 }
 
 type BTree struct {
+	Args     format.PostgresBTreeIn
 	PageSize uint64
 	page     *postgres.HeapPage
 }
 
 func decodeBTreePages(btree *BTree, d *decode.D) {
-	for i := 0; ; i++ {
-
+	for i := btree.Args.Page; ; i++ {
 		page := &postgres.HeapPage{}
 		if btree.page != nil {
 			// use prev page
diff --git a/format/postgres/common/pg_heap/postgres/pg_heap.go b/format/postgres/common/pg_heap/postgres/pg_heap.go
index 36239dda0..2bfc7f3fb 100644
--- a/format/postgres/common/pg_heap/postgres/pg_heap.go
+++ b/format/postgres/common/pg_heap/postgres/pg_heap.go
@@ -144,7 +144,7 @@ func Decode(heap *Heap, d *decode.D) any {
 }
 
 func decodeHeapPages(heap *Heap, d *decode.D) {
-	blockNumber := uint32(heap.Args.PageNumber + heap.Args.SegmentNumber*common.RelSegSize)
+	blockNumber := uint32(heap.Args.Page + heap.Args.Segment*common.RelSegSize)
 	count := int64(0)
 	for {
 		if end, _ := d.TryEnd(); end {
diff --git a/format/postgres/pg_btree.go b/format/postgres/pg_btree.go
index 5e510c65c..3812867e2 100644
--- a/format/postgres/pg_btree.go
+++ b/format/postgres/pg_btree.go
@@ -18,13 +18,19 @@ func init() {
 		Name:        format.PG_BTREE,
 		Description: "PostgreSQL btree index file",
 		DecodeFn:    decodePgBTree,
-		RootArray:   true,
-		RootName:    "pages",
+		DecodeInArg: format.PostgresBTreeIn{
+			Page: 0,
+		},
+		RootArray: true,
+		RootName:  "pages",
 	})
 	interp.RegisterFS(pgBTreeFS)
 }
 
 func decodePgBTree(d *decode.D, in any) any {
-	d.Endian = decode.LittleEndian
-	return postgres.DecodePgBTree(d)
+	pgIn, ok := in.(format.PostgresBTreeIn)
+	if !ok {
+		d.Fatalf("DecodeInArg must be PostgresBTreeIn!\n")
+	}
+	return postgres.DecodePgBTree(d, pgIn)
 }
diff --git a/format/postgres/pg_heap.go b/format/postgres/pg_heap.go
index 885bfe545..23d71f755 100644
--- a/format/postgres/pg_heap.go
+++ b/format/postgres/pg_heap.go
@@ -23,9 +23,9 @@ func init() {
 		Description: "PostgreSQL heap file",
 		DecodeFn:    decodePgheap,
 		DecodeInArg: format.PostgresHeapIn{
-			Flavour:       PG_FLAVOUR_POSTGRES14,
-			PageNumber:    0,
-			SegmentNumber: 0,
+			Flavour: PG_FLAVOUR_POSTGRES14,
+			Page:    0,
+			Segment: 0,
 		},
 		RootArray: true,
 		RootName:  "pages",