-- A lock taken while a sift_col build scans the table (or while the aggregate scan reads rows again from the heap) may reload the index's
-- relcache entry, which frees the state the index caches in it; what the build or the scan kept across tuples was then read as freed memory
-- (a segmentation fault, in a parallel worker too). sift.test_col_rel_inval invalidates the entry at every tuple, after it overwrote the
-- state the entry held, so a caller that kept it reads garbage.
CREATE EXTENSION pg_sift;
CREATE EXTENSION pg_sift_test;

CREATE FUNCTION rb_norm(d jsonb) RETURNS jsonb LANGUAGE sql IMMUTABLE PARALLEL SAFE AS $$
	SELECT CASE WHEN jsonb_typeof(d) <> 'object' THEN NULL
		ELSE jsonb_set(d, '{parts}', (SELECT coalesce(jsonb_agg(o ORDER BY ord), '[]') FROM jsonb_array_elements(d -> 'parts') WITH ORDINALITY t(o, ord)))
	END $$;

-- 1 800 rows, 3 per page: 600 pages, 3 superblocks of 256 pages; a few documents hold a numeric string in the float8 path
CREATE TABLE rb (id int, g int, doc jsonb, pad text) WITH (fillfactor = 10, autovacuum_enabled = off);
INSERT INTO rb
SELECT i, i % 7, jsonb_build_object('parts', jsonb_build_array(jsonb_build_object('c', 'a', 'score', i % 10), jsonb_build_object('c', 'b', 'score', (i % 7) + 0.5))),
	   repeat('x', 100)
FROM generate_series(1, 1800) i;
UPDATE rb SET doc = '{"parts": [{"c": "a", "score": "8.5"}, {"c": "b", "score": 2}]}' WHERE id % 50 = 5;
VACUUM ANALYZE rb;
SELECT pg_relation_size('rb') / 8192 >= 2 * 256 AS two_superblocks;

CREATE FUNCTION col_verified(idx regclass) RETURNS text LANGUAGE plpgsql AS $$
BEGIN PERFORM sift.col_verify(idx, true); RETURN 'ok'; END $$;

SET maintenance_work_mem = '128MB';
SET min_parallel_table_scan_size = 0;

-- The build, serial and with two workers: the index is complete and verifies against the heap
SET sift.test_col_rel_inval = on;
SET max_parallel_maintenance_workers = 0;
CREATE INDEX rb_ser ON rb USING sift_col (id, g, rb_norm(doc)) WITH (paths = 'rb_norm.parts[]{c text, score float8}');
SELECT current_setting('sift.test_col_workers') AS workers, sum(rows) = (SELECT count(*) FROM rb) AS all_rows, col_verified('rb_ser') AS verified
FROM sift.col_segments('rb_ser');
SET max_parallel_maintenance_workers = 2;
CREATE INDEX rb_par ON rb USING sift_col (id, g, rb_norm(doc)) WITH (paths = 'rb_norm.parts[]{c text, score float8}');
SELECT current_setting('sift.test_col_workers')::int > 0 AS parallel, sum(rows) = (SELECT count(*) FROM rb) AS all_rows, col_verified('rb_par') AS verified
FROM sift.col_segments('rb_par');
RESET sift.test_col_rel_inval;
DROP INDEX rb_ser, rb_par;

-- The aggregate scan reads the documents with a numeric string again from the heap: the answer is the normal plan's
SET max_parallel_workers_per_gather = 0;
SET seq_page_cost = 1000;
SET random_page_cost = 1000;
CREATE INDEX rb_col ON rb USING sift_col (id, g, doc) WITH (paths = 'doc.parts[]{c text, score float8}');
VACUUM ANALYZE rb;

CREATE FUNCTION rb_q(agg bool) RETURNS text LANGUAGE plpgsql AS $$
DECLARE r text;
BEGIN
	PERFORM set_config('sift.enable_agg', CASE WHEN agg THEN 'on' ELSE 'off' END, true);
	EXECUTE 'SELECT coalesce(jsonb_agg(to_jsonb(t) ORDER BY to_jsonb(t)::text), ''[]'')::text FROM (
		SELECT o.c, count(*), sum(o.score) FROM rb CROSS JOIN LATERAL jsonb_to_recordset(doc -> ''parts'') AS o(c text, score float8) GROUP BY o.c) t' INTO r;
	RETURN r;
END $$;

CREATE FUNCTION rb_pushed() RETURNS bool LANGUAGE plpgsql AS $$
DECLARE l text;
BEGIN
	PERFORM set_config('sift.enable_agg', 'on', true);
	FOR l IN EXECUTE 'EXPLAIN (COSTS OFF) SELECT o.c, count(*), sum(o.score) FROM rb CROSS JOIN LATERAL jsonb_to_recordset(doc -> ''parts'') AS o(c text, score float8) GROUP BY o.c' LOOP
		IF l ~ 'Custom Scan \(SiftAgg\)' THEN RETURN true; END IF;
	END LOOP;
	RETURN false;
END $$;

SELECT rb_pushed() AS pushed;
SELECT rb_q(true) AS answer \gset
SELECT rb_q(false) = :'answer' AS same_without_the_invalidation;
SET sift.test_col_rel_inval = on;
SELECT rb_q(true) = :'answer' AS same_with_it;
RESET sift.test_col_rel_inval;

-- A row that holds a value the index cannot store (a string in a float8 path) flags the leaf in the directory, which reloads the entry:
-- the insert went on to read and write the hint it kept there. The rows are all in the index afterwards
CREATE TABLE rw (id int, doc jsonb) WITH (autovacuum_enabled = off);
CREATE INDEX rw_col ON rw USING sift_col (id, doc) WITH (paths = 'doc.parts[]{c text, score float8}');
SET sift.test_col_rel_inval = on;
INSERT INTO rw SELECT i, jsonb_build_object('parts', jsonb_build_array(jsonb_build_object('c', 'a', 'score', CASE WHEN i % 5 = 0 THEN to_jsonb('8.5'::text) ELSE to_jsonb(i) END)))
FROM generate_series(1, 50) i;
INSERT INTO rw SELECT i, jsonb_build_object('parts', jsonb_build_array(jsonb_build_object('c', 'b', 'score', i))) FROM generate_series(51, 100) i;
RESET sift.test_col_rel_inval;
SELECT sum(rows) = (SELECT count(*) FROM rw) AS all_rows, col_verified('rw_col') AS verified FROM sift.col_segments('rw_col');
DROP TABLE rw;

-- The merge of segments cut by a key kept the leaves of the cut, in the cached state, across the units it wrote
CREATE TABLE rm (id int, grp int, v int) WITH (autovacuum_enabled = off);
INSERT INTO rm SELECT i, i % 7, i FROM generate_series(1, 20000) i;
CREATE INDEX rm_col ON rm USING sift_col (grp, v, id) WITH (key_cut = grp);
INSERT INTO rm SELECT i, i % 7, i FROM generate_series(20001, 30000) i;
SELECT sift.col_promote('rm_col') > 0 AS promoted;
SET sift.test_col_rel_inval = on;
SELECT sift.col_merge('rm_col', true) AS merged;
RESET sift.test_col_rel_inval;
SELECT count(*) AS segments, sum(rows) = (SELECT count(*) FROM rm) AS all_rows, col_verified('rm_col') AS verified FROM sift.col_segments('rm_col');
DROP TABLE rm;

DROP TABLE rb;
DROP FUNCTION rb_norm(jsonb);
DROP EXTENSION pg_sift_test;
DROP EXTENSION pg_sift;
DROP FUNCTION col_verified(regclass), rb_pushed(), rb_q(boolean);