Remove undefined behavior when predicting with invalid category value

src/compiler/ast_native.cc

@@ -642,6 +642,11 @@ class ASTNativeCompilerImpl {
           "split_index"_a = node->split_index,
           "right_categories_flag"_a = right_categories_flag);
       }
+
+      oss << fmt::format(
+          "((data[{split_index}].fvalue >= 0) && !isinf(data[{split_index}].fvalue) &&"
+          "(fabsf(data[{split_index}].fvalue) <= (float)(1U << FLT_MANT_DIG)) && (",
+          "split_index"_a = node->split_index);
       oss << "(tmp >= 0 && tmp < 64 && (( (uint64_t)"
           << bitmap[0] << "U >> tmp) & 1) )";
       for (size_t i = 1; i < bitmap.size(); ++i) {
@@ -650,7 +655,7 @@ class ASTNativeCompilerImpl {
             << " && (( (uint64_t)" << bitmap[i]
             << "U >> (tmp - " << (i * 64) << ") ) & 1) )";
       }
-      oss << ")";
+      oss << ")))";
       result = oss.str();
     }
     return result;

src/compiler/native/header_template.h

@@ -29,6 +29,7 @@ R"TREELITETEMPLATE(
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include <float.h>
 #include <math.h>
 #include <stdint.h>
 

src/gtil/predict.cc

@@ -14,6 +14,7 @@
 #include <vector>
 #include <cmath>
 #include <cstddef>
+#include <cfloat>
 #include "./pred_transform.h"
 
 namespace {
@@ -49,10 +50,16 @@ inline int NextNodeCategorical(float fvalue, const std::vector<uint32_t>& matchi
   if (std::isnan(fvalue)) {
     return default_child;
   }
-  const auto category_value = static_cast<uint32_t>(fvalue);
-  const bool is_matching_category = (
-      std::find(matching_categories.begin(), matching_categories.end(), category_value)
-      != matching_categories.end());
+  bool is_matching_category;
+  float max_representable_int = static_cast<float>(uint32_t(1) << FLT_MANT_DIG);
+  if (fvalue < 0 || std::fabs(fvalue) > max_representable_int || std::isinf(fvalue)) {
+    is_matching_category = false;
+  } else {
+    const auto category_value = static_cast<uint32_t>(fvalue);
+    is_matching_category = (
+        std::find(matching_categories.begin(), matching_categories.end(), category_value)
+        != matching_categories.end());
+  }
   if (categories_list_right_child) {
     return is_matching_category ? right_child : left_child;
   } else {

tests/python/test_invalid_categocial_input.py

@@ -0,0 +1,56 @@
+# -*- coding: utf-8 -*-
+# pylint: disable=missing-function-docstring
+"""Test whether Treelite handles invalid category values correctly"""
+import os
+
+import treelite
+import treelite_runtime
+import numpy as np
+import pytest
+from treelite.contrib import _libext
+
+from .util import os_compatible_toolchains
+
+@pytest.fixture(name='toy_model')
+def toy_model_fixture():
+    builder = treelite.ModelBuilder(num_feature=2)
+    tree = treelite.ModelBuilder.Tree()
+    tree[0].set_categorical_test_node(feature_id=1, left_categories=[0], default_left=True,
+                                      left_child_key=1, right_child_key=2)
+    tree[1].set_leaf_node(-1.0)
+    tree[2].set_leaf_node(1.0)
+    tree[0].set_root()
+    builder.append(tree)
+    model = builder.commit()
+
+    return model
+
+@pytest.fixture(name='test_data')
+def test_data_fixture():
+    categorical_column = np.array([-1, -0.6, -0.5, 0, 0.3, 0.7, 1, np.nan, np.inf, 1e10, -1e10],
+                                  dtype=np.float32)
+    dummy_column = np.zeros(categorical_column.shape[0], dtype=np.float32)
+    return np.column_stack((dummy_column, categorical_column))
+
+@pytest.fixture(name='ref_pred')
+def ref_pred_fixture():
+    # Negative inputs are mapped to the right child node
+    # 0.3 and 0.7 are mapped to the left child node, since they get rounded toward the zero.
+    # Missing value gets mapped to the left child node, since default_left=True
+    # inf, 1e10, and -1e10 don't match any element of left_categories, so they get mapped to the
+    # right child.
+    return np.array([1, 1, 1, -1, -1, -1, 1, -1, 1, 1, 1], dtype=np.float32)
+
+def test_gtil(toy_model, test_data, ref_pred):
+    pred = treelite.gtil.predict(toy_model, test_data)
+    np.testing.assert_equal(pred, ref_pred)
+
+def test_treelite_compiled(tmpdir, toy_model, test_data, ref_pred):
+    libpath = os.path.join(tmpdir, 'mylib' + _libext())
+    toolchain = os_compatible_toolchains()[0]
+    toy_model.export_lib(toolchain=toolchain, libpath=libpath)
+
+    predictor = treelite_runtime.Predictor(libpath=libpath)
+    dmat = treelite_runtime.DMatrix(test_data)
+    pred = predictor.predict(dmat)
+    np.testing.assert_equal(pred, ref_pred)