Merge pull request duckdb#11073 from carlopi/add-ipv6-inet

Add ipv6 inet + minor fixes

extension/inet/include/inet_functions.hpp

@@ -19,7 +19,9 @@ struct INetFunctions {
 	static bool CastINETToVarchar(Vector &source, Vector &result, idx_t count, CastParameters &parameters);
 
 	static void Host(DataChunk &args, ExpressionState &state, Vector &result);
+	static void Family(DataChunk &args, ExpressionState &state, Vector &result);
 	static void Subtract(DataChunk &args, ExpressionState &state, Vector &result);
+	static void Add(DataChunk &args, ExpressionState &state, Vector &result);
 };
 
 } // namespace duckdb

extension/inet/include/ipaddress.hpp

@@ -11,6 +11,7 @@
 #include "duckdb/common/common.hpp"
 #include "duckdb/common/types.hpp"
 #include "duckdb/common/types/string_type.hpp"
+#include "duckdb/common/uhugeint.hpp"
 
 namespace duckdb {
 struct CastParameters;
@@ -20,18 +21,21 @@ enum class IPAddressType : uint8_t { IP_ADDRESS_INVALID = 0, IP_ADDRESS_V4 = 1,
 class IPAddress {
 public:
 	constexpr static const int32_t IPV4_DEFAULT_MASK = 32;
+	constexpr static const int32_t IPV6_DEFAULT_MASK = 128;
+	constexpr static const int32_t IPV6_QUIBBLE_BITS = 16;
+	constexpr static const int32_t IPV6_NUM_QUIBBLE = 8;
 
 public:
 	IPAddress();
-	IPAddress(IPAddressType type, hugeint_t address, uint16_t mask);
+	IPAddress(IPAddressType type, uhugeint_t address, uint16_t mask);
 
 	IPAddressType type;
-	hugeint_t address;
+	uhugeint_t address;
 	uint16_t mask;
 
 public:
 	static IPAddress FromIPv4(int32_t address, uint16_t mask);
-	static IPAddress FromIPv6(hugeint_t address, uint16_t mask);
+	static IPAddress FromIPv6(uhugeint_t address, uint16_t mask);
 	static bool TryParse(string_t input, IPAddress &result, CastParameters &parameters);
 	static IPAddress FromString(string_t input);
 

extension/inet/inet_extension.cpp

@@ -21,6 +21,8 @@ void InetExtension::Load(DuckDB &db) {
 	// add the "inet" type
 	child_list_t<LogicalType> children;
 	children.push_back(make_pair("ip_type", LogicalType::UTINYINT));
+	// The address type would ideally be UHUGEINT, but the initial version was HUGEINT
+	// so maintain backwards-compatibility with db written with older versions.
 	children.push_back(make_pair("address", LogicalType::HUGEINT));
 	children.push_back(make_pair("mask", LogicalType::USMALLINT));
 	auto inet_type = LogicalType::STRUCT(std::move(children));
@@ -36,10 +38,15 @@ void InetExtension::Load(DuckDB &db) {
 	// add inet functions
 	ExtensionUtil::RegisterFunction(*db.instance,
 	                                ScalarFunction("host", {inet_type}, LogicalType::VARCHAR, INetFunctions::Host));
+	ExtensionUtil::RegisterFunction(
+	    *db.instance, ScalarFunction("family", {inet_type}, LogicalType::UTINYINT, INetFunctions::Family));
 
 	// Add - function with ALTER_ON_CONFLICT
-	ScalarFunction substract_fun("-", {inet_type, LogicalType::BIGINT}, inet_type, INetFunctions::Subtract);
+	ScalarFunction substract_fun("-", {inet_type, LogicalType::HUGEINT}, inet_type, INetFunctions::Subtract);
 	ExtensionUtil::AddFunctionOverload(*db.instance, substract_fun);
+
+	ScalarFunction add_fun("+", {inet_type, LogicalType::HUGEINT}, inet_type, INetFunctions::Add);
+	ExtensionUtil::AddFunctionOverload(*db.instance, add_fun);
 }
 
 std::string InetExtension::Name() {

extension/inet/inet_functions.cpp

@@ -3,13 +3,43 @@
 #include "duckdb/common/string_util.hpp"
 #include "duckdb/common/pair.hpp"
 #include "duckdb/common/operator/cast_operators.hpp"
+#include "duckdb/common/operator/subtract.hpp"
+#include "duckdb/common/operator/add.hpp"
 #include "duckdb/common/types/cast_helpers.hpp"
 #include "duckdb/common/vector_operations/generic_executor.hpp"
 
 namespace duckdb {
 
+// While the address field is better represented as a uhugeint_t, the original
+// implementation used hugeint_t, so to maintain backward-compatibility it will
+// continue to be stored as signed. However, operations on the address values
+// will use the unsigned variant, so use the functions below to convert to/from
+// the compatible representation.
 using INET_TYPE = StructTypeTernary<uint8_t, hugeint_t, uint16_t>;
 
+static uhugeint_t FromCompatAddr(hugeint_t compat_addr, IPAddressType addr_type) {
+	uhugeint_t retval = static_cast<uhugeint_t>(compat_addr);
+	// Only flip the bit for order on IPv6 addresses. It can never be set in IPv4
+	if (addr_type == IPAddressType::IP_ADDRESS_V6) {
+		// The top bit is flipped when storing as the signed hugeint so that sorting
+		// works correctly. Flip it back here to have a proper unsigned value.
+		retval.upper ^= (uint64_t(1) << 63);
+	}
+
+	return retval;
+}
+
+static hugeint_t ToCompatAddr(uhugeint_t new_addr, IPAddressType addr_type) {
+	if (addr_type == IPAddressType::IP_ADDRESS_V6) {
+		// Flip the top bit when storing as a signed hugeint_t so that sorting
+		// works correctly.
+		new_addr.upper ^= (uint64_t(1) << 63);
+	}
+	// Don't need to flip the bit for IPv4, and the original IPv4 only
+	// implementation didn't do the flipping, so maintain compatibility.
+	return static_cast<hugeint_t>(new_addr);
+}
+
 bool INetFunctions::CastVarcharToINET(Vector &source, Vector &result, idx_t count, CastParameters &parameters) {
 	auto constant = source.GetVectorType() == VectorType::CONSTANT_VECTOR;
 
@@ -37,7 +67,7 @@ bool INetFunctions::CastVarcharToINET(Vector &source, Vector &result, idx_t coun
 			continue;
 		}
 		ip_type[i] = uint8_t(inet.type);
-		address_data[i] = inet.address;
+		address_data[i] = ToCompatAddr(inet.address, inet.type);
 		mask_data[i] = inet.mask;
 	}
 	if (constant) {
@@ -48,7 +78,9 @@ bool INetFunctions::CastVarcharToINET(Vector &source, Vector &result, idx_t coun
 
 bool INetFunctions::CastINETToVarchar(Vector &source, Vector &result, idx_t count, CastParameters &parameters) {
 	GenericExecutor::ExecuteUnary<INET_TYPE, PrimitiveType<string_t>>(source, result, count, [&](INET_TYPE input) {
-		IPAddress inet(IPAddressType(input.a_val), input.b_val, input.c_val);
+		auto addr_type = IPAddressType(input.a_val);
+		auto unsigned_addr = FromCompatAddr(input.b_val, addr_type);
+		IPAddress inet(addr_type, unsigned_addr, input.c_val);
 		auto str = inet.ToString();
 		return StringVector::AddString(result, str);
 	});
@@ -58,25 +90,68 @@ bool INetFunctions::CastINETToVarchar(Vector &source, Vector &result, idx_t coun
 void INetFunctions::Host(DataChunk &args, ExpressionState &state, Vector &result) {
 	GenericExecutor::ExecuteUnary<INET_TYPE, PrimitiveType<string_t>>(
 	    args.data[0], result, args.size(), [&](INET_TYPE input) {
-		    IPAddress inet(IPAddressType(input.a_val), input.b_val, IPAddress::IPV4_DEFAULT_MASK);
+		    auto inetType = IPAddressType(input.a_val);
+		    auto mask =
+		        inetType == IPAddressType::IP_ADDRESS_V4 ? IPAddress::IPV4_DEFAULT_MASK : IPAddress::IPV6_DEFAULT_MASK;
+		    auto unsigned_addr = FromCompatAddr(input.b_val, inetType);
+		    IPAddress inet(inetType, unsigned_addr, mask);
 		    auto str = inet.ToString();
 		    return StringVector::AddString(result, str);
 	    });
 }
 
-void INetFunctions::Subtract(DataChunk &args, ExpressionState &state, Vector &result) {
-	GenericExecutor::ExecuteBinary<INET_TYPE, PrimitiveType<int32_t>, INET_TYPE>(
-	    args.data[0], args.data[1], result, args.size(), [&](INET_TYPE ip, PrimitiveType<int32_t> val) {
-		    auto new_address = ip.b_val - val.val;
-		    if (new_address < 0) {
-			    throw NotImplementedException("Out of range!?");
-		    }
-		    INET_TYPE result;
-		    result.a_val = ip.a_val;
-		    result.b_val = new_address;
-		    result.c_val = ip.c_val;
-		    return result;
+void INetFunctions::Family(DataChunk &args, ExpressionState &state, Vector &result) {
+	GenericExecutor::ExecuteUnary<INET_TYPE, PrimitiveType<uint8_t>>(
+	    args.data[0], result, args.size(), [&](INET_TYPE input) {
+		    auto inetType = IPAddressType(input.a_val);
+		    return inetType == IPAddressType::IP_ADDRESS_V4 ? 4 : 6;
 	    });
 }
 
+// The signed hugeint_t value cannot extend through the full IPv6 range in one
+// operation, but it is the largest native signed type available and should be
+// appropriate for most realistic operations. Using the signed type will make
+// the add/subtract SQL interface the most natural.
+static INET_TYPE AddImplementation(INET_TYPE ip, hugeint_t val) {
+	if (val == 0) {
+		return ip;
+	}
+
+	INET_TYPE result;
+	auto addr_type = IPAddressType(ip.a_val);
+	uhugeint_t address_in = FromCompatAddr(ip.b_val, addr_type);
+	uhugeint_t address_out;
+	result.a_val = ip.a_val;
+	result.c_val = ip.c_val;
+
+	// Use the overflow checking operators to ensure well-defined behavior. The
+	// operators must operate on the same type signedness, so convert the operand as
+	// necessary, and choose between add/subtraction operations.
+	if (val > 0) {
+		address_out = AddOperatorOverflowCheck::Operation<uhugeint_t, uhugeint_t, uhugeint_t>(address_in, val);
+	} else {
+		address_out = SubtractOperatorOverflowCheck::Operation<uhugeint_t, uhugeint_t, uhugeint_t>(address_in, -val);
+	}
+
+	if (addr_type == IPAddressType::IP_ADDRESS_V4 && address_out >= (uhugeint_t(0xffffffff))) {
+		throw OutOfRangeException("Cannot add %s to %s.", val, IPAddress(addr_type, address_in, ip.c_val).ToString());
+	}
+
+	result.b_val = ToCompatAddr(address_out, addr_type);
+
+	return result;
+}
+
+void INetFunctions::Subtract(DataChunk &args, ExpressionState &state, Vector &result) {
+	GenericExecutor::ExecuteBinary<INET_TYPE, PrimitiveType<hugeint_t>, INET_TYPE>(
+	    args.data[0], args.data[1], result, args.size(),
+	    [&](INET_TYPE ip, PrimitiveType<hugeint_t> val) { return AddImplementation(ip, -val.val); });
+}
+
+void INetFunctions::Add(DataChunk &args, ExpressionState &state, Vector &result) {
+	GenericExecutor::ExecuteBinary<INET_TYPE, PrimitiveType<hugeint_t>, INET_TYPE>(
+	    args.data[0], args.data[1], result, args.size(),
+	    [&](INET_TYPE ip, PrimitiveType<hugeint_t> val) { return AddImplementation(ip, val.val); });
+}
+
 } // namespace duckdb