use tuples to pass the pattern as enum

elixir-nx · Oct 8, 2023 · 3f663c4 · 3f663c4
1 parent b3bfd26
commit 3f663c4
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 10 deletions.
diff --git a/lib/tokenizers/pre_tokenizer.ex b/lib/tokenizers/pre_tokenizer.ex
@@ -145,15 +145,15 @@ defmodule Tokenizers.PreTokenizer do
     * `:invert` - whether to invert the split or not. Defaults to `false`
 
   """
-  @spec split(String.t() | Regex.t(), split_delimiter_behaviour(), keyword()) :: t()
+  @spec split(String.t() | {:string, String.t()}| {:regex, String.t()} , split_delimiter_behaviour(), keyword()) :: t()
   def split(pattern, behavior, opts \\ [])
 
   def split(pattern, behavior, opts) when is_binary(pattern) do
-    Tokenizers.Native.pre_tokenizers_split(pattern, behavior, opts)
+    split({:string, pattern}, behavior, opts)
   end
 
-  def split(%Regex{} = pattern, behavior, opts) do
-    split(Regex.source(pattern), behavior, Keyword.put(opts, :use_regex, true))
+  def split(pattern, behavior, opts) do
+    Tokenizers.Native.pre_tokenizers_split(pattern, behavior, opts)
   end
 
 

diff --git a/native/ex_tokenizers/src/pre_tokenizers.rs b/native/ex_tokenizers/src/pre_tokenizers.rs
@@ -240,25 +240,32 @@ impl From<SplitDelimiterBehavior> for tokenizers::SplitDelimiterBehavior {
 #[derive(NifTaggedEnum)]
 pub enum SplitOption {
     Invert(bool),
-    UseRegex(bool)
+}
+
+#[derive(NifTaggedEnum)]
+pub enum LocalSplitPattern {
+    String(String),
+    Regex(String)
 }
 
 #[rustler::nif]
 pub fn pre_tokenizers_split(
-    pattern: String,
+    pattern: LocalSplitPattern,
     behavior: SplitDelimiterBehavior,
     options: Vec<SplitOption>,
 ) -> Result<ExTokenizersPreTokenizer, rustler::Error> {
     struct Opts {
         invert: bool,
     }
     let mut opts = Opts { invert: false };
-    let mut final_pattern = SplitPattern::String(String::from(""));
+    let final_pattern = match pattern {
+        LocalSplitPattern::String(pattern) => SplitPattern::String(pattern),
+        LocalSplitPattern::Regex(pattern) => SplitPattern::Regex(pattern),
+    };
+
     for option in options {
         match option {
             SplitOption::Invert(invert) => opts.invert = invert,
-            SplitOption::UseRegex(true) => final_pattern = SplitPattern::Regex(pattern.to_owned()),
-            SplitOption::UseRegex(false) => final_pattern = SplitPattern::String(pattern.to_owned()),
         }
     }
 

diff --git a/test/tokenizers/pre_tokenizer_test.exs b/test/tokenizers/pre_tokenizer_test.exs
@@ -25,7 +25,7 @@ defmodule Tokenizers.PreTokenizerTest do
 
     test "accepts regular expressions" do
       assert %Tokenizers.PreTokenizer{} =
-               Tokenizers.PreTokenizer.split(~r/.*/, :removed)
+               Tokenizers.PreTokenizer.split({:regex, ~S/.*/}, :removed)
     end
   end