From 83d67d5a29289b15bfda977b0fd532bb53829e94 Mon Sep 17 00:00:00 2001
From: gwenn <gtreguier@gmail.com>
Date: Sat, 20 Jan 2024 19:16:48 +0100
Subject: [PATCH] Test sub-type

---
 src/functions.rs    | 44 ++++++++++++++++++++++++++++++++++++++++++--
 src/types/to_sql.rs |  7 +++++++
 2 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/src/functions.rs b/src/functions.rs
index 00ac69f..0b8c2b3 100644
--- a/src/functions.rs
+++ b/src/functions.rs
@@ -263,11 +263,18 @@ pub trait SqlFnOutput {
 }
 
 impl<T: ToSql> SqlFnOutput for T {
+    #[inline]
     fn to_sql(&self) -> Result<(ToSqlOutput<'_>, SubType)> {
         ToSql::to_sql(self).map(|o| (o, None))
     }
 }
 
+impl<T: ToSql> SqlFnOutput for (T, SubType) {
+    fn to_sql(&self) -> Result<(ToSqlOutput<'_>, SubType)> {
+        ToSql::to_sql(&self.0).map(|o| (o, self.1))
+    }
+}
+
 unsafe fn sql_result<T: SqlFnOutput>(ctx: *mut sqlite3_context, r: Result<T>) {
     let t = r.as_ref().map(SqlFnOutput::to_sql);
 
@@ -802,8 +809,8 @@ mod test {
 
     #[cfg(feature = "window")]
     use crate::functions::WindowAggregate;
-    use crate::functions::{Aggregate, Context, FunctionFlags};
-    use crate::{Connection, Error, Result};
+    use crate::functions::{Aggregate, Context, FunctionFlags, SubType};
+    use crate::{Connection, Error, Result, ValueRef};
 
     fn half(ctx: &Context<'_>) -> Result<c_double> {
         assert_eq!(ctx.len(), 1, "called with unexpected number of arguments");
@@ -1080,4 +1087,37 @@ mod test {
         assert_eq!(expected, results);
         Ok(())
     }
+
+    #[test]
+    fn test_sub_type() -> Result<()> {
+        fn test_getsubtype(ctx: &Context<'_>) -> Result<i32> {
+            Ok(ctx.get_subtype(0) as i32)
+        }
+        fn test_setsubtype<'a>(ctx: &'a Context<'_>) -> Result<(ValueRef<'a>, SubType)> {
+            use std::os::raw::c_uint;
+            let value = ctx.get_raw(0);
+            let sub_type = ctx.get::<c_uint>(1)?;
+            Ok((value, Some(sub_type)))
+        }
+        let db = Connection::open_in_memory()?;
+        db.create_scalar_function(
+            "test_getsubtype",
+            1,
+            FunctionFlags::SQLITE_UTF8,
+            test_getsubtype,
+        )?;
+        db.create_scalar_function(
+            "test_setsubtype",
+            2,
+            FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_RESULT_SUBTYPE,
+            test_setsubtype,
+        )?;
+        let result: i32 = db.one_column("SELECT test_getsubtype('hello');")?;
+        assert_eq!(0, result);
+
+        let result: i32 = db.one_column("SELECT test_getsubtype(test_setsubtype('hello',123));")?;
+        assert_eq!(123, result);
+
+        Ok(())
+    }
 }
diff --git a/src/types/to_sql.rs b/src/types/to_sql.rs
index 29e63ab..855d339 100644
--- a/src/types/to_sql.rs
+++ b/src/types/to_sql.rs
@@ -292,6 +292,13 @@ impl ToSql for Value {
     }
 }
 
+impl<'a> ToSql for ValueRef<'a> {
+    #[inline]
+    fn to_sql(&self) -> Result<ToSqlOutput<'_>> {
+        Ok(ToSqlOutput::Borrowed(*self))
+    }
+}
+
 impl<T: ToSql> ToSql for Option<T> {
     #[inline]
     fn to_sql(&self) -> Result<ToSqlOutput<'_>> {