Magic Nat

TODO.md

@@ -3,15 +3,20 @@
 
 Syntax -> Parser.Impl ?
 
+- [ ] Eq Nat is not being identified as tail recursive...
 - [x] fix string highlighting
 - [x] implement tail call optimization
-- [ ] implement magic nat (need primitive `+`, '-', and `==` in `CompileExp`)
+- [x] implement magic nat 
 - [ ] drop erased args on types and top level functions
 - [ ] can I do some inlining without blowing up code size?
 - [ ] use hint table for auto solving. (I think walking the `toList` is a big chunk of performance in `Elab.newt`.)
 - [x] implement string enum (or number, but I'm using strings for tags at the moment)
 - [x] use monaco input method instead of lean's
 - [x] `Def` is shadowed between Types and Syntax (TCon vs DCon), detect this
+- [ ] constructor magic for Bool?
+  - We'd have to make assumptions about order.
+  - we could special case some translations
+  - extra code.
 - [ ] review pattern matching. goal is to have a sane context on the other end. secondary goal - bring it closer to the paper.
   - Two issues
     - I'm rewriting stuff in the context, leaving it in a bad state (forward references). I think I can avoid this.
@@ -105,8 +110,8 @@ More comments in code! This is getting big enough that I need to re-find my bear
   - [x] records
   - [ ] record sugar? (detailed above)
   - [x] where
-  - [ ] add namespaces
-  - [ ] magic nat?
+- [x] add namespaces
+- [x] magic nat?
 - [x] rework `unify` case tree
   - Idris needs help with the case tree to keep code size down, do it in stages, one dcon at a time.
   - I'm not sure it can go a few steps deep and have a default hanging off the side, so we may need to put the default case in another function ourselves.

newt/Prelude.newt

@@ -314,23 +314,8 @@ pfunc split uses (Nil _::_) : String -> String -> List String := `(s, by) => {
 pfunc slen : String -> Int := `s => s.length`
 pfunc sindex : String -> Int -> Char := `(s,i) => s[i]`
 
--- TODO represent Nat as number at runtime
-pfunc natToInt : Nat -> Int := `(n) => {
-  let rval = 0
-  while (n.tag === 'S') {
-    n = n.h0
-    rval++
-  }
-  return rval
-}`
-
-pfunc intToNat uses (Z S) : Int -> Nat := `(n) => {
-  let rval = Prelude_Z
-  for (;n>0;n--) rval = Prelude_S(rval);
-  return rval;
-}`
-
-
+pfunc natToInt : Nat -> Int := `(n) => n`
+pfunc intToNat : Int -> Nat := `(n) => n>0?n:0`
 
 pfunc fastConcat uses (listToArray) : List String → String := `(xs) => Prelude_listToArray(null, xs).join('')`
 pfunc replicate uses (natToInt) : Nat -> Char → String := `(n,c) => c.repeat(Prelude_natToInt(n))`
@@ -796,6 +781,10 @@ find : ∀ a. (a → Bool) → List a → Maybe a
 find f Nil = Nothing
 find f (x :: xs) = if f x then Just x else find f xs
 
+any : ∀ a. (a → Bool) → List a → Bool
+any f Nil = False
+any f (x :: xs) = if f x then True else any f xs
+
 -- TODO this would be faster, but less pure as a primitive
 -- fastConcat might be a good compromise
 joinBy : String → List String → String

src/Lib/Compile.newt

@@ -195,14 +195,14 @@ termToJS {e} env (CCase t alts) f =
 
   where
     termToJSAlt : JSEnv -> String -> CAlt -> JAlt
-    termToJSAlt env nm (CConAlt name args u) = JConAlt name (termToJS (mkEnv nm 0 env args) u f)
+    termToJSAlt env nm (CConAlt name info args u) = JConAlt name (termToJS (mkEnv nm 0 env args) u f)
     -- intentionally reusing scrutinee name here
     termToJSAlt env nm (CDefAlt u) = JDefAlt (termToJS  (env) u f)
     termToJSAlt env nm (CLitAlt lit u) = JLitAlt (litToJS lit) (termToJS env u f)
 
     maybeCaseStmt : JSEnv -> String -> List CAlt -> JSStmt e
     -- If there is a single alt, assume it matched
-    maybeCaseStmt env nm ((CConAlt _ args u) :: Nil) = (termToJS (mkEnv nm 0 env args) u f)
+    maybeCaseStmt env nm ((CConAlt _ info args u) :: Nil) = (termToJS (mkEnv nm 0 env args) u f)
     maybeCaseStmt env nm alts@(CLitAlt _ _  :: _) =
         (JCase (Var nm) (map (termToJSAlt env nm) alts))
     maybeCaseStmt env nm alts =
@@ -237,6 +237,8 @@ jsIdent id = if elem id jsKeywords then text ("$" ++ id) else text $ pack $ fix
       else
         '$' :: (toHex (cast x)) ++ fix xs
 
+
+
 stmtToDoc : ∀ e. JSStmt e -> Doc
 
 
@@ -260,7 +262,7 @@ expToDoc (JLam nms body) = text "(" <+> commaSep (map jsIdent nms) <+> text ") =
 expToDoc JUndefined = text "null"
 expToDoc (Index obj ix) = expToDoc obj ++ text "(" ++ expToDoc ix ++ text " :: Nil)"
 expToDoc (Dot obj nm) = expToDoc obj ++ text "." ++ jsIdent nm
-expToDoc (JPrimOp op t u) = text "(" ++ expToDoc t <+> text op <+> expToDoc u ++ text ")"
+expToDoc (JPrimOp op t u) = parens 0 1 (expToDoc t) <+> text op <+> parens 0 1 (expToDoc u)
 
 caseBody : ∀ e. JSStmt e -> Doc
 caseBody stmt@(JReturn x) = nest 2 (line ++ stmtToDoc stmt)
@@ -341,7 +343,7 @@ sortedNames : SortedMap QName CExp → QName → List QName
 sortedNames defs qn = go Nil Nil qn
   where
     getBody : CAlt → CExp
-    getBody (CConAlt _ _ t) = t
+    getBody (CConAlt _ _ _ t) = t
     getBody (CLitAlt _ t) = t
     getBody (CDefAlt t) = t
 

src/Lib/CompileExp.newt

@@ -16,10 +16,15 @@ import Lib.Util
 import Lib.Ref2
 import Data.SortedMap
 
+-- REVIEW Separate pass for constructor magic?
+-- ConCase SuccCon will be replaced by Default CLet,
+-- but we would need to fix up zero, since we collapse extra constructors into a default case.
+-- But should be ok becaon CLitAlt doesn't bind.
+
 CExp : U
 
 data CAlt : U where
-  CConAlt : String -> List String -> CExp -> CAlt
+  CConAlt : String → ConInfo → List String → CExp → CAlt
   -- REVIEW keep var name?
   CDefAlt : CExp -> CAlt
   -- literal
@@ -109,7 +114,10 @@ compileTerm t@(Ref fc nm@(QN _ tag)) = do
     (S Z) => pure $ CRef nm
     Z =>
       case the (Maybe Def) $ lookupMap' nm defs of
-        (Just (DCon EnumCon _ _)) => pure $ CLit $ LString tag
+        Just (DCon EnumCon _ _) => pure $ CLit $ LString tag
+        Just (DCon ZeroCon _ _) => pure $ CLit $ LInt 0
+        Just (DCon SuccCon _ _) =>
+          pure $ CLam "x" $ CPrimOp "+" (CLit $ LInt 1) (CBnd 0)
         _ => pure $ CRef nm
     _ => apply (CRef nm) Nil Lin arity
 
@@ -121,20 +129,28 @@ compileTerm tm@(App _ _ _) = case funArgs tm of
         -- info (getFC tm) "Compiling an unsolved meta \{show tm}"
         -- pure $ CApp (CRef "Meta\{show k}") Nil 0
   (t@(Ref fc nm), args) => do
+        defs <- getRef Defs
         args' <- traverse compileTerm args
         arity <- arityForName fc nm
-        apply (CRef nm) args' Lin arity
+        case the (Maybe Def) $ lookupMap' nm defs of
+          Just (DCon SuccCon _ _) => applySucc args'
+          _ => apply (CRef nm) args' Lin arity
   (t, args) => do
         debug $ \ _ => "apply other \{render 90 $ pprint Nil t}"
         t' <- compileTerm t
         args' <- traverse compileTerm args
         apply t' args' Lin Z
+  where
+    applySucc : List CExp → M CExp
+    applySucc Nil = pure $ CLam "x" $ CPrimOp "+" (CLit $ LInt 1) (CBnd 0)
+    applySucc (t :: Nil) = pure $ CPrimOp "+" (CLit $ LInt 1) t
+    applySucc _ = error emptyFC "overapplied Succ \{show tm}"
 compileTerm (UU _) = pure $ CRef (QN Nil "U")
 compileTerm (Pi _ nm icit rig t u) = do
   t' <- compileTerm t
   u' <- compileTerm u
   pure $ CApp (CRef (QN Nil "PiType")) (t' :: CLam nm u' :: Nil) 0
-compileTerm (Case _ t alts) = do
+compileTerm (Case fc t alts) = do
   t' <- compileTerm t
   alts' <- for alts $ \case
     CaseDefault tm => CDefAlt <$> compileTerm tm
@@ -143,10 +159,58 @@ compileTerm (Case _ t alts) = do
       defs <- getRef Defs
       def <- lookupDef emptyFC qn
       case def of
-        DCon EnumCon _ _ => CLitAlt (LString nm) <$> compileTerm tm
-        _ => CConAlt nm args <$> compileTerm tm
+        DCon info _ _ => CConAlt nm info args <$> compileTerm tm
+        _ => error fc "\{show nm} is not constructor"
+
     CaseLit lit tm => CLitAlt lit <$> compileTerm tm
-  pure $ CCase t' alts'
+  pure $ CCase t' $ fancyCons t' alts'
+  where
+    numAltP : CAlt → Bool
+    numAltP (CConAlt _ SuccCon _ _) = True
+    numAltP (CConAlt _ ZeroCon _ _) = True
+    numAltP _ = False
+
+    enumAlt : CAlt → CAlt
+    enumAlt (CConAlt nm EnumCon args tm) = CLitAlt (LString nm) tm
+    enumAlt alt = alt
+
+    isInfo : ConInfo → CAlt → Bool
+    isInfo needle (CConAlt _ info _ _)  = needle == info
+    isInfo _ _ = False
+
+    isDef : CAlt → Bool
+    isDef (CDefAlt _) = True
+    isDef _ = False
+
+    getBody : CAlt → CExp
+    getBody (CConAlt _ _ _ t) = t
+    getBody (CLitAlt _ t) = t
+    getBody (CDefAlt t) = t
+
+    doNumCon : CExp → List CAlt → List CAlt
+    doNumCon sc alts =
+      let zeroAlt = case find (isInfo ZeroCon) alts of
+            Just (CConAlt _ _ _ tm) => CLitAlt (LInt 0) tm :: Nil
+            Just tm => fatalError "ERROR zeroAlt mismatch \{debugStr tm}"
+            _ => case find isDef alts of
+                  Just (CDefAlt tm) => CLitAlt (LInt 0) tm :: Nil
+                  -- This happens if the zero alt is impossible
+                  _ => Nil
+      in
+      let succAlt = case find (isInfo SuccCon) alts of
+            Just (CConAlt _ _ _ tm) => CDefAlt (CLet "x" (CPrimOp "-" sc (CLit $ LInt 1)) tm) :: Nil
+            Just tm => fatalError "ERROR succAlt mismatch \{debugStr tm}"
+            _ => case find isDef alts of
+              Just alt => alt :: Nil
+              _ => Nil
+      in zeroAlt ++ succAlt
+
+    fancyCons : CExp → List CAlt → List CAlt
+    fancyCons sc alts =
+      if any numAltP alts
+        then doNumCon sc alts
+        else map enumAlt alts
+
 compileTerm (Lit _ lit) = pure $ CLit lit
 compileTerm (Let _ nm t u) = do
   t' <- compileTerm t

src/Lib/TCO.newt

@@ -23,7 +23,7 @@ tailNames (CApp (CRef name) args 0) = name :: Nil
 tailNames (CCase _ alts) = join $ map altTailNames alts
   where
     altTailNames : CAlt → List QName
-    altTailNames (CConAlt _ _ exp) = tailNames exp
+    altTailNames (CConAlt _ _ _ exp) = tailNames exp
     altTailNames (CDefAlt exp) = tailNames exp
     altTailNames (CLitAlt _ exp) = tailNames exp
 tailNames (CLet _ _ t) = tailNames t
@@ -54,7 +54,7 @@ rewriteTailCalls nms tm = case tm of
   tm => CConstr "return" (tm :: Nil)
   where
       rewriteAlt : CAlt -> CAlt
-      rewriteAlt (CConAlt nm args t) = CConAlt nm args $ rewriteTailCalls nms t
+      rewriteAlt (CConAlt nm info args t) = CConAlt nm info args $ rewriteTailCalls nms t
       rewriteAlt (CDefAlt t) = CDefAlt $ rewriteTailCalls nms t
       rewriteAlt (CLitAlt lit t) = CLitAlt lit $ rewriteTailCalls nms t
 
@@ -85,7 +85,7 @@ doOptimize fns = do
     mkRecName (QN ns nm :: _) = pure $ QN ns "REC_\{nm}"
 
     mkAlt : List QName → (QName × List Name × CExp) -> CAlt
-    mkAlt nms (qn, args, tm) = CConAlt (show qn) args (rewriteTailCalls nms tm)
+    mkAlt nms (qn, args, tm) = CConAlt (show qn) NormalCon args (rewriteTailCalls nms tm)
 
     splitFun : (QName × CExp) → M (QName × List Name × CExp)
     splitFun (qn, CFun args body) = pure (qn, args, body)

src/Lib/Types.newt

@@ -310,6 +310,13 @@ record MetaContext where
 
 data ConInfo = NormalCon | SuccCon | ZeroCon | EnumCon
 
+instance Eq ConInfo where
+  NormalCon == NormalCon = True
+  SuccCon == SuccCon = True
+  ZeroCon == ZeroCon = True
+  EnumCon == EnumCon = True
+  _ == _ = False
+
 instance Show ConInfo where
   show NormalCon = ""
   show SuccCon = "[S]"