improved number operations

Author: tjdevries

Cargo.lock

@@ -8,6 +8,7 @@ edition = "2021"
 [dependencies]
 lexer = { path = "../vim9-lexer", package = "vim9-lexer" }
 parser = { path = "../vim9-parser", package = "vim9-parser" }
+vimfuncs = { path = "../vimfuncs" }
 
 anyhow = "1.0.62"
 stylua = "0.14.3"

crates/vim9-gen/Cargo.toml

@@ -789,15 +789,14 @@ impl Generate for IfCommand {
         //
         // This would allow us to remove the NVIM9.bool condition (assuming the type
         // system is correct in vim9script)
-        //
-        // let condition = match guess_type_of_expr(state, &self.condition).inner {
-        //     Type::Bool => condition,
-        //     _ => format!("NVIM9.bool({condition})"),
-        // };
+        let condition = match guess_type_of_expr(state, &self.condition) {
+            Type::Bool => condition,
+            _ => format!("NVIM9.bool({condition})"),
+        };
 
         format!(
             r#"
-if NVIM9.bool({condition}) then
+if {condition} then
   {}
 {}
 {}
@@ -881,36 +880,65 @@ fn identifier_list(state: &mut State, unpacked: &UnpackIdentifier) -> String {
         .collect()
 }
 
+fn vim_to_type(ret: &vimfuncs::FuncReturnType) -> Type {
+    match ret {
+        vimfuncs::FuncReturnType::Any => Type::Any,
+        vimfuncs::FuncReturnType::Bool => Type::Bool,
+        vimfuncs::FuncReturnType::Float => Type::Float,
+        vimfuncs::FuncReturnType::Number => Type::Number,
+        vimfuncs::FuncReturnType::NumberBool => Type::BoolOrNumber,
+        vimfuncs::FuncReturnType::String => Type::String,
+        // vimfuncs::FuncReturnType::Void => Type::Any,
+        // vimfuncs::FuncReturnType::Dict(_) => Type,
+        // vimfuncs::FuncReturnType::List(_) => todo!(),
+        // vimfuncs::FuncReturnType::Func(_) => todo!(),
+        _ => Type::Any,
+    }
+}
+
 fn guess_type_of_expr(state: &State, expr: &Expression) -> Type {
     match expr {
         Expression::Number(_) => Type::Number,
         Expression::String(_) => Type::String,
         Expression::Boolean(_) => Type::Bool,
-        Expression::Call(c) => match c.name() {
-            Some(ident) => match ident {
-                // TODO: Use our very cool new generated stuff here
-                Identifier::Raw(raw) => match raw.name.as_str() {
-                    "charcol" => Type::Number,
-                    _ => Type::Any,
-                },
+        Expression::Call(c) => {
+            match c.name() {
+                Some(ident) => {
+                    match ident {
+                        // TODO: Use our very cool new generated stuff here
+                        Identifier::Raw(raw) => {
+                            match vimfuncs::get_func_info(&raw.name) {
+                                Some(info) => vim_to_type(&info.return_type),
+                                _ => Type::Any,
+                            }
+                        }
+                        _ => Type::Any,
+                    }
+                }
                 _ => Type::Any,
-            },
-            _ => Type::Any,
-        },
+            }
+        }
         Expression::Infix(infix) => {
             if infix.operator.is_comparison() {
                 return Type::Bool;
             }
 
+            let left_ty = guess_type_of_expr(state, infix.left.as_ref());
+            let right_ty = guess_type_of_expr(state, infix.right.as_ref());
+
             if matches!(infix.operator, Operator::And | Operator::Or) {
-                if guess_type_of_expr(state, infix.left.as_ref()) == Type::Bool
-                    && guess_type_of_expr(state, infix.right.as_ref())
-                        == Type::Bool
-                {
+                if left_ty == Type::Bool && right_ty == Type::Bool {
                     return Type::Bool;
                 }
             }
 
+            if left_ty == Type::Number
+                && right_ty == Type::Number
+                && infix.operator.is_math()
+            {
+                return dbg!(Type::Number);
+            }
+
             Type::Any
         }
         Expression::Grouped(g) => guess_type_of_expr(state, g.expr.as_ref()),
@@ -1348,6 +1376,11 @@ impl Generate for InfixExpression {
         if let Some(literal) = self.operator.literal() {
             match ty {
                 Bool => return format!("{left} {literal} {right}",),
+                Number => {
+                    if self.operator.is_math() {
+                        return format!("{left} {literal} {right}");
+                    }
+                }
                 _ => {}
             }
         }

crates/vim9-gen/src/lib.rs

@@ -35,7 +35,7 @@ describe("filename", function()
     for _, y in NVIM9.iter({ 1, 2, 3 }) do
       x = NVIM9.ops["Plus"](x, y)
 
-      if NVIM9.bool(y == 2) then
+      if y == 2 then
         break
       end
     end
@@ -56,11 +56,11 @@ describe("filename", function()
     local body = function(_, y)
       x = NVIM9.ops["Plus"](x, y)
 
-      if NVIM9.bool(y == 2) then
+      if y == 2 then
         return NVIM9.ITER_CONTINUE
       end
 
-      if NVIM9.bool(y == 3) then
+      if y == 3 then
         return NVIM9.ITER_BREAK
       end
 
@@ -90,7 +90,7 @@ describe("filename", function()
     local x = 0
 
     local body = function(_, y)
-      if NVIM9.bool(y == 2) then
+      if y == 2 then
         return NVIM9.ITER_CONTINUE
       end
 
@@ -124,11 +124,11 @@ describe("filename", function()
       local x = 0
 
       local body = function(_, y)
-        if NVIM9.bool(y == 2) then
+        if y == 2 then
           return NVIM9.ITER_RETURN, 5
         end
 
-        if NVIM9.bool(y == 10) then
+        if y == 10 then
           return NVIM9.ITER_CONTINUE
         end
 

crates/vim9-gen/testdata/output/busted_loops.lua

@@ -1,13 +1,13 @@
 ---
 source: crates/vim9-gen/src/lib.rs
-assertion_line: 1361
-expression: "generate(contents, false)"
+assertion_line: 1586
+expression: "generate(contents, false).unwrap()"
 ---
 local NVIM9 = require("vim9script")
 local __VIM9_MODULE = {}
 -- vim9script
 
-vim.g["my_var"] = NVIM9.ops["Plus"](1, 2)
+vim.g["my_var"] = 1 + 2
 print(vim.g["my_var"])
 return __VIM9_MODULE
 

crates/vim9-gen/testdata/output/vim9_gen__test__assign.snap

@@ -1,6 +1,6 @@
 ---
 source: crates/vim9-gen/src/lib.rs
-assertion_line: 1553
+assertion_line: 1577
 expression: "generate(contents, false).unwrap()"
 ---
 local NVIM9 = require("vim9script")
@@ -83,21 +83,21 @@ Qf_filter = function(qf, searchpat, bang)
 
   firstchar = NVIM9.index(searchpat, 0)
   lastchar = NVIM9.slice(searchpat, -1, nil)
-  if NVIM9.bool(firstchar == lastchar and (firstchar == "/" or firstchar == '"' or firstchar == "'")) then
+  if firstchar == lastchar and (firstchar == "/" or firstchar == '"' or firstchar == "'") then
     pat = NVIM9.slice(searchpat, 1, -2)
-    if NVIM9.bool(pat == "") then
+    if pat == "" then
       -- # Use the last search pattern
       pat = vim.fn.getreg("/")
     end
   else
     pat = searchpat
   end
 
-  if NVIM9.bool(pat == "") then
+  if pat == "" then
     return
   end
 
-  if NVIM9.bool(bang == "!") then
+  if bang == "!" then
     Cond = function(_, val)
       return NVIM9.ops["NotRegexpMatches"](val["text"], pat)
         and NVIM9.ops["NotRegexpMatches"](NVIM9.fn["bufname"](val["bufnr"]), pat)

crates/vim9-gen/testdata/output/vim9_gen__test__cfilter.snap

@@ -1,13 +1,13 @@
 ---
 source: crates/vim9-gen/src/lib.rs
-assertion_line: 1359
-expression: "generate(contents, false)"
+assertion_line: 1584
+expression: "generate(contents, false).unwrap()"
 ---
 local NVIM9 = require("vim9script")
 local __VIM9_MODULE = {}
 -- vim9script
 
-local x = NVIM9.ops["Plus"](1, 2)
+local x = 1 + 2
 print(x)
 return __VIM9_MODULE
 

crates/vim9-gen/testdata/output/vim9_gen__test__expr.snap

@@ -1331,6 +1331,11 @@ impl Operator {
         )
     }
 
+    pub fn is_math(&self) -> bool {
+        use Operator::*;
+        matches!(self, Plus | Minus | Multiply | Divide)
+    }
+
     pub fn literal(&self) -> Option<&'static str> {
         Some(match self {
             Operator::Plus => "+",

crates/vim9-parser/src/lib.rs

@@ -6,7 +6,8 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+phf = "0.11.1"
 
-[[bin]]
-name = "main"
-path = "src/main.rs"
+[build-dependencies]
+anyhow = "1.0.66"
+phf_codegen = "0.11.1"