From 7625cc2728dbe38973bb0c2e46cd9e0deaaa8607 Mon Sep 17 00:00:00 2001
From: Shane O'Brien <shane.obrien@circuithub.com>
Date: Thu, 19 Oct 2023 18:07:38 +0100
Subject: [PATCH] Add `index` and `index1` for indexing arrays (ListTable and
 NonEmptyTable)

---
 .../20231019_180613_shane.obrien_index.md     |  3 +++
 src/Rel8.hs                                   | 20 ++++++++-----------
 src/Rel8/Expr/List.hs                         | 18 +++++++++++++----
 src/Rel8/Expr/NonEmpty.hs                     | 20 ++++++++++++++-----
 src/Rel8/Table/List.hs                        | 13 +++++++++++-
 src/Rel8/Table/NonEmpty.hs                    | 15 +++++++++++++-
 src/Rel8/Type/Array.hs                        | 18 ++++++++++++-----
 7 files changed, 79 insertions(+), 28 deletions(-)
 create mode 100644 changelog.d/20231019_180613_shane.obrien_index.md

diff --git a/changelog.d/20231019_180613_shane.obrien_index.md b/changelog.d/20231019_180613_shane.obrien_index.md
new file mode 100644
index 00000000..11354a47
--- /dev/null
+++ b/changelog.d/20231019_180613_shane.obrien_index.md
@@ -0,0 +1,3 @@
+### Added
+
+- Added `index`, `index1`, `indexExpr`, and `index1Expr` functions for extracting individual elements from `ListTable`s and `NonEmptyTable`s.
diff --git a/src/Rel8.hs b/src/Rel8.hs
index e0a259df..6695d4a7 100644
--- a/src/Rel8.hs
+++ b/src/Rel8.hs
@@ -101,12 +101,10 @@ module Rel8
   , manyExpr
   , catListTable
   , catList
-  , head
-  , headExpr
-  , last
-  , lastExpr
-  , length
-  , lengthExpr
+  , head, headExpr
+  , index, indexExpr
+  , last, lastExpr
+  , length, lengthExpr
 
     -- ** @NonEmptyTable@
   , NonEmptyTable
@@ -116,12 +114,10 @@ module Rel8
   , someExpr
   , catNonEmptyTable
   , catNonEmpty
-  , head1
-  , head1Expr
-  , last1
-  , last1Expr
-  , length1
-  , length1Expr
+  , head1, head1Expr
+  , index1, index1Expr
+  , last1, last1Expr
+  , length1, length1Expr
 
     -- ** @NullTable@
   , NullTable
diff --git a/src/Rel8/Expr/List.hs b/src/Rel8/Expr/List.hs
index 3158a792..1797b5f8 100644
--- a/src/Rel8/Expr/List.hs
+++ b/src/Rel8/Expr/List.hs
@@ -3,8 +3,10 @@
 
 module Rel8.Expr.List (
   headExpr,
+  indexExpr,
   lastExpr,
   sheadExpr,
+  sindexExpr,
   slastExpr,
   lengthExpr,
 ) where
@@ -15,7 +17,7 @@ import Prelude
 
 -- rel8
 import Rel8.Expr (Expr)
-import Rel8.Expr.Opaleye (fromPrimExpr, toPrimExpr)
+import Rel8.Expr.Opaleye (fromPrimExpr, mapPrimExpr, toPrimExpr)
 import Rel8.Schema.Null (Nullify, Sql, Unnullify)
 import Rel8.Type (DBType, typeInformation)
 import Rel8.Type.Information (TypeInformation)
@@ -26,17 +28,25 @@ headExpr :: Sql DBType a => Expr [a] -> Expr (Nullify a)
 headExpr = sheadExpr typeInformation
 
 
+indexExpr :: Sql DBType a => Expr Int32 -> Expr [a] -> Expr (Nullify a)
+indexExpr = sindexExpr typeInformation
+
+
 lastExpr :: Sql DBType a => Expr [a] -> Expr (Nullify a)
 lastExpr = slastExpr typeInformation
 
 
 sheadExpr :: TypeInformation (Unnullify a) -> Expr [a] -> Expr (Nullify a)
-sheadExpr info = fromPrimExpr . Prim.head info . toPrimExpr
+sheadExpr info = mapPrimExpr (Prim.head info)
+
+
+sindexExpr :: TypeInformation (Unnullify a) -> Expr Int32 -> Expr [a] -> Expr (Nullify a)
+sindexExpr info i = mapPrimExpr (Prim.index info (toPrimExpr i))
 
 
 slastExpr :: TypeInformation (Unnullify a) -> Expr [a] -> Expr (Nullify a)
-slastExpr info = fromPrimExpr . Prim.last info . toPrimExpr
+slastExpr info = mapPrimExpr (Prim.last info)
 
 
 lengthExpr :: Expr [a] -> Expr Int32
-lengthExpr = fromPrimExpr . Prim.length . toPrimExpr
+lengthExpr = mapPrimExpr (Prim.length)
diff --git a/src/Rel8/Expr/NonEmpty.hs b/src/Rel8/Expr/NonEmpty.hs
index 82bfd83d..70ff5e77 100644
--- a/src/Rel8/Expr/NonEmpty.hs
+++ b/src/Rel8/Expr/NonEmpty.hs
@@ -3,8 +3,10 @@
 
 module Rel8.Expr.NonEmpty (
   head1Expr,
+  index1Expr,
   last1Expr,
   shead1Expr,
+  sindex1Expr,
   slast1Expr,
   length1Expr,
 ) where
@@ -16,8 +18,8 @@ import Prelude
 
 -- rel8
 import Rel8.Expr (Expr)
-import Rel8.Expr.Opaleye (fromPrimExpr, toPrimExpr)
-import Rel8.Schema.Null (Sql, Unnullify)
+import Rel8.Expr.Opaleye (fromPrimExpr, mapPrimExpr, toPrimExpr)
+import Rel8.Schema.Null (Nullify, Sql, Unnullify)
 import Rel8.Type (DBType, typeInformation)
 import Rel8.Type.Information (TypeInformation)
 import qualified Rel8.Type.Array as Prim
@@ -27,17 +29,25 @@ head1Expr :: Sql DBType a => Expr (NonEmpty a) -> Expr a
 head1Expr = shead1Expr typeInformation
 
 
+index1Expr :: Sql DBType a => Expr Int32 -> Expr (NonEmpty a) -> Expr (Nullify a)
+index1Expr = sindex1Expr typeInformation
+
+
 last1Expr :: Sql DBType a => Expr (NonEmpty a) -> Expr a
 last1Expr = slast1Expr typeInformation
 
 
 shead1Expr :: TypeInformation (Unnullify a) -> Expr (NonEmpty a) -> Expr a
-shead1Expr info = fromPrimExpr . Prim.head info . toPrimExpr
+shead1Expr info = mapPrimExpr (Prim.head info)
+
+
+sindex1Expr :: TypeInformation (Unnullify a) -> Expr Int32 -> Expr (NonEmpty a) -> Expr (Nullify a)
+sindex1Expr info i = mapPrimExpr (Prim.index info (toPrimExpr i))
 
 
 slast1Expr :: TypeInformation (Unnullify a) -> Expr (NonEmpty a) -> Expr a
-slast1Expr info = fromPrimExpr . Prim.last info . toPrimExpr
+slast1Expr info = mapPrimExpr (Prim.last info)
 
 
 length1Expr :: Expr (NonEmpty a) -> Expr Int32
-length1Expr = fromPrimExpr . Prim.length . toPrimExpr
+length1Expr = mapPrimExpr (Prim.length)
diff --git a/src/Rel8/Table/List.hs b/src/Rel8/Table/List.hs
index 57e31054..3c5b7202 100644
--- a/src/Rel8/Table/List.hs
+++ b/src/Rel8/Table/List.hs
@@ -16,6 +16,7 @@ module Rel8.Table.List
   , listTable
   , nameListTable
   , head
+  , index
   , last
   , length
   )
@@ -30,7 +31,7 @@ import Prelude hiding (head, last, length)
 -- rel8
 import Rel8.Expr ( Expr )
 import Rel8.Expr.Array ( sappend, sempty, slistOf )
-import Rel8.Expr.List (lengthExpr, sheadExpr, slastExpr)
+import Rel8.Expr.List (lengthExpr, sheadExpr, sindexExpr, slastExpr)
 import Rel8.Schema.Dict ( Dict( Dict ) )
 import Rel8.Schema.HTable.List ( HListTable )
 import Rel8.Schema.HTable.Vectorize
@@ -166,6 +167,16 @@ head =
   toColumns
 
 
+-- | @'index' i as@ extracts a single element from @as@, returning
+-- 'Rel8.nullTable' if @i@ is out of range. Note that although PostgreSQL
+-- array indexes are 1-based (by default), this function is always 0-based.
+index :: Table Expr a => Expr Int32 -> ListTable Expr a -> NullTable Expr a
+index i =
+  fromColumns .
+  hnullify (\Spec {info} -> sindexExpr info i) .
+  toColumns
+
+
 -- | Get the last element of a 'ListTable' (or 'Rel8.nullTable' if empty).
 last :: Table Expr a => ListTable Expr a -> NullTable Expr a
 last =
diff --git a/src/Rel8/Table/NonEmpty.hs b/src/Rel8/Table/NonEmpty.hs
index c8fc849c..a37e0397 100644
--- a/src/Rel8/Table/NonEmpty.hs
+++ b/src/Rel8/Table/NonEmpty.hs
@@ -16,6 +16,7 @@ module Rel8.Table.NonEmpty
   , nonEmptyTable
   , nameNonEmptyTable
   , head1
+  , index1
   , last1
   , length1
   )
@@ -31,11 +32,12 @@ import Prelude hiding ( id )
 -- rel8
 import Rel8.Expr ( Expr )
 import Rel8.Expr.Array ( sappend1, snonEmptyOf )
-import Rel8.Expr.NonEmpty (length1Expr, shead1Expr, slast1Expr)
+import Rel8.Expr.NonEmpty (length1Expr, shead1Expr, sindex1Expr, slast1Expr)
 import Rel8.Schema.Dict ( Dict( Dict ) )
 import Rel8.Schema.HTable.NonEmpty ( HNonEmptyTable )
 import Rel8.Schema.HTable.Vectorize
   ( hvectorize, hunvectorize
+  , hnullify
   , happend
   , hproject, hcolumn
   , First (..)
@@ -52,6 +54,7 @@ import Rel8.Table
   )
 import Rel8.Table.Alternative ( AltTable, (<|>:) )
 import Rel8.Table.Eq ( EqTable, eqTable )
+import Rel8.Table.Null (NullTable)
 import Rel8.Table.Ord ( OrdTable, ordTable )
 import Rel8.Table.Projection
   ( Projectable, Projecting, Projection, project, apply
@@ -159,6 +162,16 @@ head1 =
   toColumns
 
 
+-- | @'index1' i as@ extracts a single element from @as@, returning
+-- 'Rel8.nullTable' if @i@ is out of range. Note that although PostgreSQL
+-- array indexes are 1-based (by default), this function is always 0-based.
+index1 :: Table Expr a => Expr Int32 -> NonEmptyTable Expr a -> NullTable Expr a
+index1 i =
+  fromColumns .
+  hnullify (\Spec {info} -> sindex1Expr info i) .
+  toColumns
+
+
 -- | Get the last element of a 'NonEmptyTable'.
 last1 :: Table Expr a => NonEmptyTable Expr a -> a
 last1 =
diff --git a/src/Rel8/Type/Array.hs b/src/Rel8/Type/Array.hs
index 304a2984..42b75a6c 100644
--- a/src/Rel8/Type/Array.hs
+++ b/src/Rel8/Type/Array.hs
@@ -10,7 +10,7 @@ module Rel8.Type.Array
   , arrayTypeName
   , listTypeInformation
   , nonEmptyTypeInformation
-  , head, last, length
+  , head, index, last, length
   )
 where
 
@@ -157,15 +157,19 @@ arrayParser = \case
 
 
 head :: TypeInformation a -> Opaleye.PrimExpr -> Opaleye.PrimExpr
-head info a = extractArrayElement info $ index (lower a) a
+head info a = extractArrayElement info $ subscript (lower a) a
 
 
 last :: TypeInformation a -> Opaleye.PrimExpr -> Opaleye.PrimExpr
-last info a = extractArrayElement info $ index (upper a) a
+last info a = extractArrayElement info $ subscript (upper a) a
 
 
-index :: Opaleye.PrimExpr -> Opaleye.PrimExpr -> Opaleye.PrimExpr
-index i a = Opaleye.ArrayIndex a i
+subscript :: Opaleye.PrimExpr -> Opaleye.PrimExpr -> Opaleye.PrimExpr
+subscript i a = Opaleye.ArrayIndex a i
+
+
+index :: TypeInformation a -> Opaleye.PrimExpr -> Opaleye.PrimExpr -> Opaleye.PrimExpr
+index info i a = extractArrayElement info $ subscript (plus (lower a) i) a
 
 
 lower :: Opaleye.PrimExpr -> Opaleye.PrimExpr
@@ -186,3 +190,7 @@ one = Opaleye.ConstExpr (Opaleye.IntegerLit 1)
 
 zero :: Opaleye.PrimExpr
 zero = Opaleye.ConstExpr (Opaleye.IntegerLit 0)
+
+
+plus :: Opaleye.PrimExpr -> Opaleye.PrimExpr -> Opaleye.PrimExpr
+plus = Opaleye.BinExpr (Opaleye.:+)
\ No newline at end of file