wip:raw-opt improvements

aslanix · aslanix · commit a2ab13751eb8 · 2025-08-01T09:25:32.000+02:00
diff --git a/compiler/src/RawOpt.hs b/compiler/src/RawOpt.hs
@@ -96,9 +96,14 @@ data PState =
              stateSubst :: Subst,
              stateChange:: ChangeFlag,
              stateRawVarTypes :: Map RawVar RawType,              -- for assertions optimizations
-             stateLValTypes :: Map VarName RawType                -- for assertions optimizations
+             stateLValTypes :: Map VarName RawType                -- for assertions optimizations             
            }
 
+
+data ReadEnv = 
+    ReadEnv { readConsts :: Map RawVar Core.Lit }
+           
+
 -- 2021-02-28; AA 
 -- As we traverse the AST we collect information about how 
 -- different bindings are used. We distinguish two different 
@@ -113,7 +118,7 @@ type Used = (Set VarName, Set RawVar)
 type ChangeFlag = Bool
 
 -- Optimization monad: keep track of used variables, to be able to eliminate unused variables.
-type Opt = RWS () Used PState
+type Opt = RWS ReadEnv Used PState
 
 class PEval a where 
     peval :: a -> Opt a 
@@ -266,13 +271,13 @@ _setLValType x t = modify (\pstate ->
 -- Partially evaluate instruction. This is called multiple times in the optimization sequence.
 -- First pass: partially evaluate functions (instructions).
 -- Removes e.g. redundant state projections state (e.g. if multiple in same block).
--- Returns 'Nothing' if the instruction is to be ommitted.
-pevalInst:: RawInst -> Opt (Maybe RawInst)
+-- 
+pevalInst:: RawInst -> Opt [RawInst]
 pevalInst i = do 
     pstate <- get
     i' <- subst i -- apply the collected substitutions
-    let _omit x = x >> (return Nothing)
-    let _keep x = x >> (return $ Just i')
+    let _omit x = x >> return []
+    let _keep x = x >> return [i']
              
     case i' of 
       AssignRaw r (ProjectState p) -> do
@@ -319,21 +324,39 @@ pevalInst i = do
         case guessType rexpr of 
             Nothing -> return ()
             Just ty -> _setRawType r ty 
-        
 
-      AssignLVal v complexExpr -> _keep $ do
-        case complexExpr of 
-            Bin op (UseNativeBinop False) r1 r2 | op `elem` [Basics.Eq, Basics.Neq]  -> 
-              _setLValType v RawBoolean
-            _ -> return ()
+      AssignLVal v complexExpr@(Bin op (UseNativeBinop False) r1 r2) 
+          | op `elem` [Basics.Eq, Basics.Neq] -> do
+                _setLValType v RawBoolean
+                a <- isSuitableForNativeEq r1
+                b <- isSuitableForNativeEq r2
+                if  a || b 
+                  then do 
+                    let VN s = v 
+                        r3 = RawVar $ s ++ "$val_opt"
+                        r4 = RawVar $ s ++ "$vlev_opt"
+                        r5 = RawVar $ s ++ "$tlev_opt"
+                    markUsed v 
+                    markUsed [r1, r2, r3, r4, r5]
+                    return $
+                      [ AssignRaw r3 (Bin op (UseNativeBinop True) r1 r2)
+                      , AssignRaw r4 (ProjectState MonPC)
+                      , AssignRaw r5 (ProjectState MonPC)
+                      , AssignLVal v (ConstructLVal r3 r4 r5)
+                      ]
+                  else 
+                    _keep $ markUsed complexExpr
+
+
+      AssignLVal v complexExpr -> 
         _keep $ markUsed complexExpr
 
       SetState p r -> _keep $ do 
         markUsed r
         monInsert p r        
       RTAssertion (AssertType r rt) -> do 
         case Map.lookup r (stateRawVarTypes pstate) of 
-          Just rt' | rt' == rt -> return Nothing
+          Just rt' | rt' == rt -> return []
           _ -> _keep $ _setRawType r rt >> markUsed r        
       -- RTAssertion (AssertEqTypes opt_ls x y) -> do
       --   let _m = stateTypes pstate
@@ -353,19 +376,34 @@ pevalInst i = do
         case (Map.lookup x _m, Map.lookup y _m) of 
           (Just t1 , Just t2) | t1 == t2 -> 
             if t1 `elem` [RawNumber, RawString]
-            then return Nothing
+            then return []
             else keep
           _ -> keep
       -- TODO track tuple length
       RTAssertion (AssertTupleLengthGreaterThan r n) -> _keep $ markUsed r
       -- TODO track record fields
       RTAssertion (AssertRecordHasField r f) -> _keep $ markUsed r
-      RTAssertion (AssertNotZero r) -> _keep $ markUsed r
+      RTAssertion (AssertNotZero r) -> do
+         renv <- ask 
+         case Map.lookup r (readConsts renv) of 
+           Just (Core.LInt x _) | x /= 0 -> return []
+           _ -> _keep $ markUsed r
       MkFunClosures ee _ -> _keep $ markUsed (snd (unzip ee)) 
       -- No applicable optimizations.
-      SetBranchFlag -> return $ Just i'
-      InvalidateSparseBit -> return $ Just i'
-                          
+      SetBranchFlag -> return [i']
+      InvalidateSparseBit -> return [i']
+
+
+isSuitableForNativeEq r = do 
+  pstate <- get
+  return $ 
+    case Map.lookup r (stateRawVarTypes pstate) of 
+      Nothing -> False 
+      Just t -> case t of 
+          RawNumber -> True 
+          RawString -> True 
+          _ -> False 
+
 
 instance PEval RawTerminator where
   peval tr = do 
@@ -488,7 +526,7 @@ instOrder ii = work [] ii
 instance PEval RawBBTree where
   peval bb@(BB insts tr) = do 
     (BB jj tr'', used) <- listen $ do 
-        ii <- Data.Maybe.catMaybes <$> mapM pevalInst insts
+        ii <- concat <$> mapM pevalInst insts
         tr' <- peval tr 
         return $ BB ii tr'
     let (insts_no_ret, set_pc_bl) = filterInstBwd (filter (isLiveInstFwd used) jj)
@@ -538,8 +576,11 @@ funopt (FunDef hfn consts bb ir) =
                        stateRawVarTypes = constTypes,
                        stateLValTypes = Map.empty
                        }
-      (bb', _, _) = runRWS (peval bb) () pstate
-      new = FunDef hfn consts' bb' ir  
+      
+      readenv = ReadEnv { readConsts = Map.fromList consts  }
+      (bb', _, (_, used_rvars)) = runRWS (peval bb) readenv pstate
+      const_used = filter (\(x,_) -> Set.member x used_rvars) consts'
+      new = FunDef hfn const_used bb' ir  
   in if bb /= bb' then funopt new else new
 
 
diff --git a/compiler/src/Stack2JS.hs b/compiler/src/Stack2JS.hs
@@ -93,7 +93,7 @@ addOneLib (LibAccess (Basics.LibName libname) varname) =
   let args = (PP.quotes.PP.text) libname <+> text "," <+> (PP.quotes. PP.text) varname
   in text "this.addLib " <+> PP.parens args
 
-addLibs xs = vcat (map addOneLib xs)
+addLibs xs = vcat $ nub (map addOneLib xs)
 
 
 data TheState = TheState { freshCounter :: Integer
@@ -312,8 +312,8 @@ instance ToJS StackInst where
 instance ToJS StackTerminator where 
   toJS = tr2js
 
-binOpToJS :: BinOp -> String
-binOpToJS = \case
+binOpToJS :: BinOp -> Raw.UseNativeBinop ->  String
+binOpToJS op (Raw.UseNativeBinop isNative) = case op of 
     -- JS binary operators (some not implemented in IR2Raw)
     Plus -> "+"
     Minus -> "-"
@@ -334,8 +334,8 @@ binOpToJS = \case
     BinZeroShiftRight -> ">>>"
     -- Functions defined in UserRuntimeZero.ts
     IntDiv -> "rt.intdiv"
-    Eq -> "rt.eq"
-    Neq -> "rt.neq"
+    Eq -> if isNative then "===" else "rt.eq"
+    Neq -> if isNative then "!==" else "rt.neq"
     Concat -> "+"
     HasField -> "rt.hasField"
     LatticeJoin -> "rt.raw_join"
@@ -414,7 +414,7 @@ ir2js (MkFunClosures envBindings funBindings) = do
                               d_b <- toJS b 
                               return $ d_b <> text ".dataLevel") ls
                let d2 = penv PP.<> text ".__dataLevel = " 
-                        <+> jsFunCall (text $ binOpToJS Basics.LatticeJoin) d3
+                        <+> jsFunCall (text $ binOpToJS Basics.LatticeJoin (Raw.UseNativeBinop False)) d3
                                 
                return $ vcat ( d1 ++ [d2])
              hsepc ls = semi $ PP.hsep (PP.punctuate (text ",") ls)
@@ -583,9 +583,9 @@ instance ToJS RawExpr where
            Raw.FieldValLev -> monStateToJs MonPC
            Raw.FieldTypLev -> monStateToJs MonPC)
       e@(ProjectLVal _ _) -> return $ ppRawExpr e
-      Bin binop _ va1 va2 -> return $
-        let text' = (text . binOpToJS) binop in
-          if isInfixBinop binop
+      Bin binop use_native va1 va2 -> return $
+        let text' = text (binOpToJS binop use_native) in
+          if isInfixBinop binop use_native
           then hsep [ ppId va1, text', ppId va2 ]
           else jsFunCall text' [ppId va1, ppId va2]
       Un op v -> return $ text (unaryOpToJS op) <> PP.parens (ppId v)
@@ -663,8 +663,8 @@ freshKontName = do
     return $ VN $  "$$$" ++ s ++ "$$$kont" ++ (show j)
 
 
-isInfixBinop :: Basics.BinOp -> Bool
-isInfixBinop = \case
+isInfixBinop :: Basics.BinOp -> Raw.UseNativeBinop -> Bool
+isInfixBinop op (Raw.UseNativeBinop use_native) = case op of 
   -- Infix
   Plus -> True
   Minus -> True
@@ -684,9 +684,10 @@ isInfixBinop = \case
   BinShiftLeft -> True
   BinShiftRight -> True
   BinZeroShiftRight -> True
+  -- Flag dependent 
+  Eq -> use_native
+  Neq -> use_native 
   -- Not infix
-  Eq -> False
-  Neq -> False
   RaisedTo -> False
   FlowsTo -> False
   IntDiv -> False
diff --git a/tests/_unautomated/loop.trp b/tests/_unautomated/loop.trp
@@ -1,6 +1,6 @@
 let
   fun loop n =
-    if n = 10 then
+    if n <> 10 then
       let val _ = printWithLabels n
       in loop (n + 1)
       end
