Better modular arithmetic optmizer

core/inference/class/class-tests.factor

@@ -13,9 +13,10 @@ system layouts vectors ;
 ! Ensure type inference works as it is supposed to by checking
 ! if various methods get inlined
 
-: inlined? ( quot word -- ? )
+: inlined? ( quot seq/word -- ? )
+    dup word? [ 1array ] when
     swap dataflow optimize
-    [ node-param eq? ] with node-exists? not ;
+    [ node-param swap member? ] with node-exists? not ;
 
 GENERIC: mynot ( x -- y )
 
@@ -323,3 +324,48 @@ cell-bits 32 = [
         ] when
     ] \ + inlined?
 ] unit-test
+
+[ f ] [
+    [
+        256 mod
+    ] { mod fixnum-mod } inlined?
+] unit-test
+
+[ f ] [
+    [
+        dup 0 >= [ 256 mod ] when
+    ] { mod fixnum-mod } inlined?
+] unit-test
+
+[ t ] [
+    [
+        { integer } declare dup 0 >= [ 256 mod ] when
+    ] { mod fixnum-mod } inlined?
+] unit-test
+
+[ t ] [
+    [
+        { integer } declare 256 rem
+    ] { mod fixnum-mod } inlined?
+] unit-test
+
+! [ t ] [
+!     [
+!         { integer } declare [ 256 mod ] map
+!     ] { mod fixnum-mod } inlined?
+! ] unit-test
+! 
+! [ t ] [
+!     [
+!         { integer } declare [ 0 >= ] map
+!     ] { >= fixnum>= } inlined?
+! ] unit-test
+
+[ t ] [
+    [
+        { integer } declare
+        dup 0 >= [
+            615949 * 797807 + 20 2^ mod dup 19 2^ -
+        ] [ dup ] if
+    ] { * + shift mod fixnum-mod fixnum* fixnum+ fixnum- } inlined?
+] unit-test

core/math/integers/integers-tests.factor

@@ -184,3 +184,10 @@ unit-test
 [ HEX: 988a259c3433f237 ] [
     B{ HEX: 37 HEX: f2 HEX: 33 HEX: 34 HEX: 9c HEX: 25 HEX: 8a HEX: 98 } byte-array>bignum
 ] unit-test
+
+[ t ] [ 256 power-of-2? ] unit-test
+[ f ] [ 123 power-of-2? ] unit-test
+
+[ f ] [ -128 power-of-2? ] unit-test
+[ f ] [ 0 power-of-2? ] unit-test
+[ t ] [ 1 power-of-2? ] unit-test

core/math/math-docs.factor

@@ -1,5 +1,5 @@
 USING: help.markup help.syntax kernel sequences quotations
-math.private math.functions ;
+math.private ;
 IN: math
 
 ARTICLE: "division-by-zero" "Division by zero"
@@ -26,17 +26,13 @@ $nl
 { $subsection < }
 { $subsection <= }
 { $subsection > }
-{ $subsection >= }
-"Inexact comparison:"
-{ $subsection ~ } ;
+{ $subsection >= } ;
 
 ARTICLE: "modular-arithmetic" "Modular arithmetic"
 { $subsection mod }
 { $subsection rem }
 { $subsection /mod }
 { $subsection /i }
-{ $subsection mod-inv }
-{ $subsection ^mod }
 { $see-also "integer-functions" } ;
 
 ARTICLE: "bitwise-arithmetic" "Bitwise arithmetic"
@@ -363,6 +359,10 @@ HELP: next-power-of-2
 { $values { "m" "a non-negative integer" } { "n" "an integer" } }
 { $description "Outputs the smallest power of 2 greater than " { $snippet "m" } ". The output value is always at least 1." } ;
 
+HELP: power-of-2?
+{ $values { "n" integer } { "?" "a boolean" } }
+{ $description "Tests if " { $snippet "n" } " is a power of 2." } ;
+
 HELP: each-integer
 { $values { "n" integer } { "quot" "a quotation with stack effect " { $snippet "( i -- )" } } }
 { $description "Applies the quotation to each integer from 0 up to " { $snippet "n" } ", excluding " { $snippet "n" } "." }

core/math/math.factor

@@ -121,7 +121,11 @@ M: float fp-nan?
 
 : next-power-of-2 ( m -- n ) 2 swap (next-power-of-2) ; foldable
 
-: align ( m w -- n ) 1- [ + ] keep bitnot bitand ; inline
+: power-of-2? ( n -- ? )
+    dup 0 < [ drop f ] [ dup 1- bitand zero? ] if ; foldable
+
+: align ( m w -- n )
+    1- [ + ] keep bitnot bitand ; inline
 
 <PRIVATE
 

core/optimizer/math/math.factor

@@ -426,18 +426,22 @@ most-negative-fixnum most-positive-fixnum [a,b]
     swap sole-consumer
     dup #call? [ node-param eq? ] [ 2drop f ] if ;
 
-: coereced-to-fixnum? ( #call -- ? )
-    \ >fixnum consumed-by? ;
+: coerced-to-fixnum? ( #call -- ? )
+    dup dup node-in-d [ node-class integer class< ] with all?
+    [ \ >fixnum consumed-by? ] [ drop f ] if ;
 
 {
     { fixnum+ [ fixnum+fast ] }
     { fixnum- [ fixnum-fast ] }
     { fixnum* [ fixnum*fast ] }
+    { + [ >r >fixnum r> >fixnum fixnum+fast ] }
+    { - [ >r >fixnum r> >fixnum fixnum-fast ] }
+    { * [ >r >fixnum r> >fixnum fixnum*fast ] }
 } [
     [
         [
             dup remove-overflow-check?
-            over coereced-to-fixnum? or
+            over coerced-to-fixnum? or
         ] ,
         [ f splice-quot ] curry ,
     ] { } make 1array define-optimizers
@@ -467,3 +471,49 @@ most-negative-fixnum most-positive-fixnum [a,b]
         [ [ fixnum-shift-fast ] f splice-quot ]
     }
 } define-optimizers
+
+: convert-rem-to-and? ( #call -- ? )
+    dup node-in-d {
+        { [ 2dup first node-class integer class< not ] [ f ] }
+        { [ 2dup second node-literal integer? not ] [ f ] }
+        { [ 2dup second node-literal power-of-2? not ] [ f ] }
+        [ t ]
+    } cond 2nip ;
+
+: convert-mod-to-and? ( #call -- ? )
+    dup dup node-in-d first node-interval 0 [a,inf] interval-subset?
+    [ convert-rem-to-and? ] [ drop f ] if ;
+
+: convert-mod-to-and ( #call -- node )
+    dup
+    dup node-in-d second node-literal 1-
+    [ nip bitand ] curry f splice-quot ;
+
+{ mod bignum-mod fixnum-mod } [
+    {
+        {
+            [ dup convert-mod-to-and? ]
+            [ convert-mod-to-and ]
+        }
+    } define-optimizers
+] each
+
+\ rem {
+    {
+        [ dup convert-rem-to-and? ]
+        [ convert-mod-to-and ]
+    }
+} define-optimizers
+
+: fixnumify-bitand? ( #call -- ? )
+    dup node-in-d second node-interval fixnum fits? ;
+
+: fixnumify-bitand ( #call -- node )
+    [ >r >fixnum r> >fixnum fixnum-bitand ] f splice-quot ;
+
+\ bitand {
+    {
+        [ dup fixnumify-bitand? ]
+        [ fixnumify-bitand ]
+    }
+} define-optimizers

extra/math/functions/functions-docs.factor

@@ -7,6 +7,9 @@ ARTICLE: "integer-functions" "Integer functions"
 { $subsection gcd }
 { $subsection log2 }
 { $subsection next-power-of-2 }
+"Modular exponentiation:"
+{ $subsection ^mod }
+{ $subsection mod-inv }
 "Tests:"
 { $subsection power-of-2? }
 { $subsection even? }
@@ -33,7 +36,9 @@ ARTICLE: "arithmetic-functions" "Arithmetic functions"
 { $subsection ceiling }
 { $subsection floor }
 { $subsection truncate }
-{ $subsection round } ;
+{ $subsection round }
+"Inexact comparison:"
+{ $subsection ~ } ;
 
 ARTICLE: "power-functions" "Powers and logarithms"
 "Squares:"
@@ -107,10 +112,6 @@ HELP: >rect
 { $values { "z" number } { "x" real } { "y" real } }
 { $description "Extracts the real and imaginary components of a complex number." } ;
 
-HELP: power-of-2?
-{ $values { "n" integer } { "?" "a boolean" } }
-{ $description "Tests if " { $snippet "n" } " is a power of 2." } ;
-
 HELP: align
 { $values { "m" integer } { "w" "a power of 2" } { "n" "an integer multiple of " { $snippet "w" } } }
 { $description "Outputs the least multiple of " { $snippet "w" } " greater than " { $snippet "m" } "." }

extra/math/functions/functions-tests.factor

@@ -81,9 +81,6 @@ IN: math.functions.tests
 [ 1/8 ] [ 2 -3 ^ ] unit-test
 [ t ] [ 1 100 shift 2 100 ^ = ] unit-test
 
-[ t ] [ 256 power-of-2? ] unit-test
-[ f ] [ 123 power-of-2? ] unit-test
-
 [ 1 ] [ 7/8 ceiling ] unit-test
 [ 2 ] [ 3/2 ceiling ] unit-test
 [ 0 ] [ -7/8 ceiling ] unit-test

extra/math/functions/functions.factor

@@ -102,9 +102,6 @@ M: real absq sq ;
         [ ~abs ]
     } cond ;
 
-: power-of-2? ( n -- ? )
-    dup 0 < [ drop f ] [ dup 1- bitand zero? ] if ; foldable
-
 : >rect ( z -- x y ) dup real-part swap imaginary-part ; inline
 
 : conjugate ( z -- z* ) >rect neg rect> ; inline