feat: Add nat hope solving

brat/Brat/Checker/Helpers.hs

@@ -272,58 +272,12 @@ vecLayers :: Modey m -> Val Z -> Checking ([(Src, NumVal (VVar Z))] -- The sizes
                                           ,CTy m Z -- The function type at the end
                                           )
 vecLayers my (TVec ty (VNum n)) = do
-  src <- mkStaticNum n
+  src <- buildNatVal n
   first ((src, n):) <$> vecLayers my ty
 vecLayers Braty (VFun Braty cty) = pure ([], cty)
 vecLayers Kerny (VFun Kerny cty) = pure ([], cty)
 vecLayers my ty = typeErr $ "Expected a " ++ showMode my ++ "function or vector of functions, got " ++ show ty
 
-mkStaticNum :: NumVal (VVar Z) -> Checking Src
-mkStaticNum n@(NumValue c gro) = do
-  (_, [], [(constSrc,_)], _) <- next "const" (Const (Num (fromIntegral c))) (S0, Some (Zy :* S0)) R0 (RPr ("value", TNat) R0)
-  src <- case gro of
-    Constant0 -> pure constSrc
-    StrictMonoFun sm -> do
-      (_, [(lhs,_),(rhs,_)], [(src,_)], _) <- next "add_const" (ArithNode Add) (S0, Some (Zy :* S0))
-                                              (RPr ("lhs", TNat) (RPr ("rhs", TNat) R0))
-                                              (RPr ("value", TNat) R0)
-      smSrc <- mkStrictMono sm
-      wire (constSrc, TNat, lhs)
-      wire (smSrc, TNat, rhs)
-      pure src
-  defineSrc src (VNum n)
-  pure src
- where
-  mkStrictMono :: StrictMono (VVar Z) -> Checking Src
-  mkStrictMono (StrictMono k mono) = do
-    (_, [], [(constSrc,_)], _) <- next "2^k" (Const (Num (2^k))) (S0, Some (Zy :* S0)) R0 (RPr ("value", TNat) R0)
-    (_, [(lhs,_),(rhs,_)], [(src,_)], _) <- next "mult_const" (ArithNode Mul) (S0, Some (Zy :* S0))
-                                            (RPr ("lhs", TNat) (RPr ("rhs", TNat) R0))
-                                            (RPr ("value", TNat) R0)
-    monoSrc <- mkMono mono
-    wire (constSrc, TNat, lhs)
-    wire (monoSrc, TNat, rhs)
-    pure src
-
-  mkMono :: Monotone (VVar Z) -> Checking Src
-  mkMono (Linear (VPar (ExEnd e))) = pure (NamedPort e "mono")
-  mkMono (Full sm) = do
-    (_, [], [(twoSrc,_)], _) <- next "2" (Const (Num 2)) (S0, Some (Zy :* S0)) R0 (RPr ("value", TNat) R0)
-    (_, [(lhs,_),(rhs,_)], [(powSrc,_)], _) <- next "2^" (ArithNode Pow) (S0, Some (Zy :* S0))
-                                               (RPr ("lhs", TNat) (RPr ("rhs", TNat) R0))
-                                               (RPr ("value", TNat) R0)
-    smSrc <- mkStrictMono sm
-    wire (twoSrc, TNat, lhs)
-    wire (smSrc, TNat, rhs)
-
-    (_, [], [(oneSrc,_)], _) <- next "1" (Const (Num 1)) (S0, Some (Zy :* S0)) R0 (RPr ("value", TNat) R0)
-    (_, [(lhs,_),(rhs,_)], [(src,_)], _) <- next "n-1" (ArithNode Sub) (S0, Some (Zy :* S0))
-                                            (RPr ("lhs", TNat) (RPr ("rhs", TNat) R0))
-                                            (RPr ("value", TNat) R0)
-    wire (powSrc, TNat, lhs)
-    wire (oneSrc, TNat, rhs)
-    pure src
-
 vectorise :: forall m. Modey m -> (Src, Val Z) -> Checking (Src, CTy m Z)
 vectorise my (src, ty) = do
   (layers, cty) <- vecLayers my ty
@@ -493,7 +447,100 @@ runArith (NumValue upl grol) Pow (NumValue upr gror)
  = pure $ NumValue (upl ^ upr) (StrictMonoFun (StrictMono (l * upr) (Full (StrictMono (k + k') mono))))
 runArith _ _ _ = Nothing
 
+buildArithOp :: ArithOp -> Checking ((Tgt, Tgt), Src)
+buildArithOp op = do
+  (_, [(lhs,_), (rhs,_)], [(out,_)], _) <- next (show op) (ArithNode op) (S0, Some (Zy :* S0)) (RPr ("lhs", TNat) (RPr ("rhs", TNat) R0)) (RPr ("value", TNat) R0)
+  pure ((lhs, rhs), out)
+
 buildConst :: SimpleTerm -> Val Z -> Checking Src
 buildConst tm ty = do
   (_, _, [(out,_)], _) <- next "buildConst" (Const tm) (S0, Some (Zy :* S0)) R0 (RPr ("value", ty) R0)
   pure out
+
+buildNum :: Integer -> Checking Src
+buildNum n = buildConst (Num (fromIntegral n)) TNat
+
+-- Generate wiring to produce a dynamic instance of the numval argument
+-- N.B. In these functions, we wire using Req, rather than the `wire` function
+-- because we don't want it to do any extra evaluation.
+buildNatVal :: NumVal (VVar Z) -> Checking Src
+buildNatVal nv@(NumValue n gro) = case n of
+  0 -> buildGro gro
+  n -> do
+    nDangling <- buildNum n
+    ((lhs,rhs),out) <- buildArithOp Add
+    src <- buildGro gro
+    req $ Wire (end nDangling, TNat, end lhs)
+    req $ Wire (end src, TNat, end rhs)
+    defineSrc out (VNum (nPlus n (nVar (VPar (toEnd src)))))
+    pure out
+ where
+  buildGro :: Fun00 (VVar Z) -> Checking Src
+  buildGro Constant0 = buildNum 0
+  buildGro (StrictMonoFun sm) = buildSM sm
+
+  buildSM :: StrictMono (VVar Z) -> Checking Src
+  buildSM (StrictMono k mono) = do
+    factor <- buildNum $ 2 ^ k
+    -- Multiply mono by 2^k; note we could avoid this if k==0
+    ((lhs,rhs),out) <- buildArithOp Mul
+    monoDangling <- buildMono mono
+    req $ Wire (end factor, TNat, end lhs)
+    req $ Wire (end monoDangling, TNat, end rhs)
+    defineSrc out (VNum (n2PowTimes k (nVar (VPar (toEnd monoDangling)))))
+    pure out
+
+  buildMono :: Monotone (VVar Z) -> Checking Src
+  buildMono (Linear (VPar (ExEnd e))) = pure $ NamedPort e "numval"
+  buildMono (Full sm) = do
+    -- Calculate 2^n as `outPlus1`
+    two <- buildNum 2
+    dangling <- buildSM sm
+    ((lhs,rhs),outPlus1) <- buildArithOp Pow
+    req $ Wire (end two, TNat, end lhs)
+    req $ Wire (end dangling, TNat, end rhs)
+    -- Then subtract 1
+    one <- buildNum 1
+    ((lhs,rhs),out) <- buildArithOp Sub
+    req $ Wire (end outPlus1, TNat, end lhs)
+    req $ Wire (end one, TNat, end rhs)
+    defineSrc out (VNum (nFull (nVar (VPar (toEnd dangling)))))
+    pure out
+  buildMono _ = err . InternalError $ "Trying to build a non-closed nat value: " ++ show nv
+
+invertNatVal :: NumVal (VVar Z) -> Checking Tgt
+invertNatVal (NumValue up gro) = case up of
+  0 -> invertGro gro
+  _ -> do
+    ((lhs,rhs),out) <- buildArithOp Sub
+    upSrc <- buildNum up
+    req $ Wire (end upSrc, TNat, end rhs)
+    tgt <- invertGro gro
+    req $ Wire (end out, TNat, end tgt)
+    defineTgt tgt (VNum (nVar (VPar (toEnd out))))
+    defineTgt lhs (VNum (nPlus up (nVar (VPar (toEnd tgt)))))
+    pure lhs
+ where
+  invertGro Constant0 = error "Invariant violated: the numval arg to invertNatVal should contain a variable"
+  invertGro (StrictMonoFun sm) = invertSM sm
+
+  invertSM (StrictMono k mono) = case k of
+    0 -> invertMono mono
+    _ -> do
+      divisor <- buildNum (2 ^ k)
+      ((lhs,rhs),out) <- buildArithOp Div
+      tgt <- invertMono mono
+      req $ Wire (end out, TNat, end tgt)
+      req $ Wire (end divisor, TNat, end rhs)
+      defineTgt tgt (VNum (nVar (VPar (toEnd out))))
+      defineTgt lhs (VNum (n2PowTimes k (nVar (VPar (toEnd tgt)))))
+      pure lhs
+
+  invertMono (Linear (VPar (InEnd e))) = pure (NamedPort e "numval")
+  invertMono (Full sm) = do
+    (_, [(llufTgt,_)], [(llufSrc,_)], _) <- next "luff" (Prim ("BRAT","lluf")) (S0, Some (Zy :* S0)) (REx ("n", Nat) R0) (REx ("n", Nat) R0)
+    tgt <- invertSM sm
+    req $ Wire (end llufSrc, TNat, end tgt)
+    defineTgt tgt (VNum (nVar (VPar (toEnd llufSrc))))
+    defineTgt llufTgt (VNum (nFull (nVar (VPar (toEnd tgt)))))
+    pure llufTgt

brat/Brat/Checker/SolveHoles.hs

@@ -1,6 +1,6 @@
 module Brat.Checker.SolveHoles (typeEq) where
 
-import Brat.Checker.Helpers (buildConst)
+import Brat.Checker.Helpers (buildConst, buildNatVal)
 import Brat.Checker.Monad
 import Brat.Checker.Types (kindForMode)
 import Brat.Error (ErrorMsg(..))
@@ -98,7 +98,7 @@ solveHope k hope v = quote Zy v >>= \v -> case doesntOccur (InEnd hope) v of
   Right () -> do
     defineEnd (InEnd hope) v
     dangling <- case (k, v) of
-      (Nat, VNum _v) -> err $ Unimplemented "Nat hope solving" []
+      (Nat, VNum v) -> buildNatVal v
       (Nat, _) -> err $ InternalError "Head of Nat wasn't a VNum"
       _ -> buildConst Unit TUnit
     req (Wire (end dangling, kindType k, hope))

brat/Brat/Checker/SolvePatterns.hs

@@ -198,6 +198,42 @@ instantiateMeta e val = do
   throwLeft (doesntOccur e val)
   defineEnd e val
 
+-- Solve a Nat kinded metavariable. Unlike `instantiateMeta`, this function also
+-- makes the dynamic wiring for a metavariable. This only needs to happen for
+-- numbers because they have nontrivial runtime behaviour.
+--
+-- We assume that the caller has done the occurs check and rules out trivial equations.
+solveNumMeta :: End -> NumVal (VVar Z) -> Checking ()
+solveNumMeta e nv = case (e, vars nv) of
+ -- Compute the thing that the rhs should be based on the src, and instantiate src to that
+ (ExEnd src,  [VPar (InEnd _tgt)]) -> do
+   -- Compute the value of the `tgt` variable from the known `src` value by inverting nv
+   tgtSrc <- invertNatVal nv
+   instantiateMeta (ExEnd src) (VNum (nVar (VPar (toEnd tgtSrc))))
+   wire (NamedPort src "", TNat, tgtSrc)
+
+ (ExEnd src, _) -> instantiateMeta (ExEnd src) (VNum nv)
+
+ -- Both targets, we need to create the thing that they both derive from
+ (InEnd bigTgt, [VPar (InEnd weeTgt)]) -> do
+   (_, [(idTgt, _)], [(idSrc, _)], _) <- anext "numval id" Id (S0, Some (Zy :* S0))
+                                         (REx ("n", Nat) R0) (REx ("n", Nat) R0)
+   defineSrc idSrc (VNum (nVar (VPar (toEnd idTgt))))
+   instantiateMeta (InEnd weeTgt) (VNum (nVar (VPar (toEnd idSrc))))
+   wire (idSrc, TNat, NamedPort weeTgt "")
+   let nv' = fmap (const (VPar (toEnd idSrc))) nv -- weeTgt is the only thing to replace
+   bigSrc <- buildNatVal nv'
+   instantiateMeta (InEnd bigTgt) (VNum nv')
+   wire (bigSrc, TNat, NamedPort bigTgt "")
+
+ -- RHS is constant or Src, wire it into tgt
+ (InEnd tgt,  _) -> do
+   src <- buildNatVal nv
+   instantiateMeta (InEnd tgt) (VNum nv)
+   wire (src, TNat, NamedPort tgt "")
+ where
+  vars :: NumVal a -> [a]
+  vars = foldMap pure
 
 -- Need to keep track of which way we're solving - which side is known/unknown
 -- Things which are dynamically unknown must be Tgts - information flows from Srcs
@@ -218,11 +254,9 @@ unifyNum (NumValue lup lgro) (NumValue rup rgro)
     lhsStrictMono (StrictMono (n - 1) mono) num
 
   lhsMono :: Monotone (VVar Z) -> NumVal (VVar Z) -> Checking ()
-  lhsMono (Linear v) num = case v of
-    VPar e -> instantiateMeta e (VNum num)
-    _ -> case num of -- our only hope is to instantiate the RHS
-      NumValue 0 (StrictMonoFun (StrictMono 0 (Linear (VPar (ExEnd e))))) -> instantiateMeta (toEnd e) (VNum (nVar v))
-      _ -> err . UnificationError $ "Couldn't instantiate variable " ++ show v
+  lhsMono (Linear v) (NumValue 0 (StrictMonoFun (StrictMono 0 (Linear v')))) | v == v' = pure ()
+  lhsMono (Linear (VPar e)) num = throwLeft (doesntOccur e (VNum num)) *>
+                                  solveNumMeta e num
   lhsMono (Full sm) (NumValue 0 (StrictMonoFun (StrictMono 0 (Full sm'))))
     = lhsStrictMono sm (NumValue 0 (StrictMonoFun sm'))
   lhsMono m@(Full _) (NumValue 0 gro) = lhsFun00 gro (NumValue 0 (StrictMonoFun (StrictMono 0 m)))
@@ -233,7 +267,7 @@ unifyNum (NumValue lup lgro) (NumValue rup rgro)
   demand0 :: NumVal (VVar Z) -> Checking ()
   demand0 (NumValue 0 Constant0) = pure ()
   demand0 n@(NumValue 0 (StrictMonoFun (StrictMono _ mono))) = case mono of
-    Linear (VPar e) -> instantiateMeta e (VNum (nConstant 0))
+    Linear (VPar e) -> solveNumMeta e (nConstant 0)
     Full sm -> demand0 (NumValue 0 (StrictMonoFun sm))
     _ -> err . UnificationError $ "Couldn't force " ++ show n ++ " to be 0"
   demand0 n = err . UnificationError $ "Couldn't force " ++ show n ++ " to be 0"
@@ -243,9 +277,23 @@ unifyNum (NumValue lup lgro) (NumValue rup rgro)
   --   2^k * x
   -- = 2^k * (y + 1)
   -- = 2^k + 2^k * y
-  demandSucc (StrictMono k (Linear (VPar (ExEnd out)))) = do
-    y <- mkPred out
-    pure $ nPlus ((2 ^ k) - 1) $ n2PowTimes k y
+  demandSucc (StrictMono k (Linear (VPar (ExEnd x)))) = do
+    (_, [(yTgt, _)], [(ySrc, _)], _) <-
+      next "yId" Id (S0, Some (Zy :* S0)) (REx ("value", Nat) R0) (REx ("value", Nat) R0)
+
+    defineSrc ySrc (VNum (nVar (VPar (toEnd yTgt))))
+    instantiateMeta (ExEnd x) (VNum (nPlus 1 (nVar (VPar (toEnd yTgt)))))
+    pure $ nPlus ((2 ^ k) - 1) $ n2PowTimes k (nVar (VPar (toEnd ySrc)))
+  --   2^k * x
+  -- = 2^k * (y + 1)
+  -- = 2^k + 2^k * y
+  -- Hence, the predecessor is (2^k - 1) + (2^k * y)
+  demandSucc (StrictMono k (Linear (VPar (InEnd x)))) = do
+    (_, [(y,_)], _, _) <- anext "y" Hypo (S0, Some (Zy :* S0)) (REx ("", Nat) R0) R0
+    yPlus1 <- invertNatVal (nPlus 1 (nVar (VPar (toEnd y))))
+    solveNumMeta (InEnd x) (nVar (VPar (toEnd yPlus1)))
+    pure $ nPlus ((2 ^ k) - 1) $ n2PowTimes k (nVar (VPar (toEnd y)))
+
   --   2^k * full(n + 1)
   -- = 2^k * (1 + 2 * full(n))
   -- = 2^k + 2^(k + 1) * full(n)
@@ -264,53 +312,39 @@ unifyNum (NumValue lup lgro) (NumValue rup rgro)
     evenGro Constant0 = pure Constant0
     evenGro (StrictMonoFun (StrictMono 0 mono)) = case mono of
       Linear (VPar (ExEnd out)) -> do
-        half <- mkHalf out
+        half <- invertNatVal (NumValue 0 (StrictMonoFun (StrictMono 1 mono)))
+        solveNumMeta (ExEnd out) (n2PowTimes 1 (nVar (VPar (toEnd half))))
         pure (StrictMonoFun (StrictMono 0 (Linear (VPar (toEnd half)))))
-      Linear _ -> err . UnificationError $ "Can't force " ++ show n ++ " to be even"
+      Linear (VPar (InEnd tgt)) -> do
+        halfTgt <- buildNatVal (NumValue 0 (StrictMonoFun (StrictMono 1 (Linear (VPar (toEnd tgt))))))
+        let half = nVar (VPar (toEnd halfTgt))
+        solveNumMeta (InEnd tgt) (n2PowTimes 1 half)
+        pure (StrictMonoFun (StrictMono 0 (Linear (VPar (toEnd halfTgt)))))
       Full sm -> StrictMonoFun sm <$ demand0 (NumValue 0 (StrictMonoFun sm))
     evenGro (StrictMonoFun (StrictMono n mono)) = pure (StrictMonoFun (StrictMono (n - 1) mono))
 
     -- Check a numval is odd, and return its rounded down half
     oddGro :: Fun00 (VVar Z) -> Checking (NumVal (VVar Z))
     oddGro (StrictMonoFun (StrictMono 0 mono)) = case mono of
-      Linear (VPar (ExEnd out)) -> mkPred out >>= demandEven
-      Linear _ -> err . UnificationError $ "Can't force " ++ show n ++ " to be even"
+      -- TODO: Why aren't we using `out`??
+      Linear (VPar (ExEnd _out)) -> do
+        -- compute (/2) . (-1)
+        doubTgt <- invertNatVal (NumValue 1 (StrictMonoFun (StrictMono 1 mono)))
+        let [VPar (InEnd halfTgt)] = foldMap pure mono
+        solveNumMeta (toEnd doubTgt) (nPlus 1 (n2PowTimes 1 (nVar (VPar (toEnd halfTgt)))))
+        pure (nVar (VPar (toEnd halfTgt)))
+      Linear (VPar (InEnd weeTgt)) -> do
+        -- compute (/2) . (-1)
+        bigTgt <- invertNatVal (NumValue 1 (StrictMonoFun (StrictMono 1 (Linear (VPar (toEnd weeTgt))))))
+        let flooredHalf = nVar (VPar (toEnd weeTgt))
+        solveNumMeta (toEnd bigTgt) (nPlus 1 (n2PowTimes 1 flooredHalf))
+        pure flooredHalf
+
       -- full(n + 1) = 1 + 2 * full(n)
       -- hence, full(n) is the rounded down half
       Full sm -> nFull <$> demandSucc sm
     oddGro _ = err . UnificationError $ "Can't force " ++ show n ++ " to be even"
 
-  -- Add dynamic logic to compute half of a variable.
-  mkHalf :: OutPort -> Checking Src
-  mkHalf out = do
-    (_, [], [(const2,_)], _) <- next "const2" (Const (Num 2)) (S0, Some (Zy :* S0))
-                                R0
-                                (RPr ("value", TNat) R0)
-    (_, [(lhs,_),(rhs,_)], [(half,_)], _) <- next "div2" (ArithNode Div) (S0, Some (Zy :* S0))
-                                             (RPr ("left", TNat) (RPr ("right", TNat) R0))
-                                             (RPr ("out", TNat) R0)
-    wire (NamedPort out "numerator", TNat, lhs)
-    wire (const2, TNat, rhs)
-    req $ Define (toEnd out) (VNum (n2PowTimes 1 (nVar (VPar (toEnd half)))))
-    pure half
-
-
-  -- Add dynamic logic to compute the predecessor of a variable, and return that
-  -- predecessor.
-  -- The variable must be a non-zero nat!!
-  mkPred :: OutPort -> Checking (NumVal (VVar Z))
-  mkPred out = do
-    (_, [], [(const1,_)], _) <- next "const1" (Const (Num 1)) (S0, Some (Zy :* S0))
-                                R0
-                                (RPr ("value", TNat) R0)
-    (_, [(lhs,_),(rhs,_)], [(pred,_)], _) <- next "minus1" (ArithNode Sub) (S0, Some (Zy :* S0))
-                                             (RPr ("left", TNat) (RPr ("right", TNat) R0))
-                                             (RPr ("out", TNat) R0)
-    wire (NamedPort out "", TNat, lhs)
-    wire (const1, TNat, rhs)
-    req $ Define (ExEnd out) (VNum (nPlus 1 (nVar (VPar (toEnd pred)))))
-    pure (nVar (VPar (toEnd pred)))
-
 -- The variable must be a non-zero nat!!
 patVal :: ValPat -> [End] -> (Val Z, [End])
 -- Nat variables will only be found in a `NumPat`, not a `ValPat`

brat/examples/infer.brat

@@ -2,7 +2,6 @@ map(X :: *, Y :: *, { X -> Y }, List(X)) -> List(Y)
 map(_, _, _, []) = []
 map(_, _, f, x ,- xs) = f(x) ,- map(!, !, f, xs)
 
--- TODO: Make BRAT solve for the # kinded args
 mapVec(X :: *, Y :: *, { X -> Y }, n :: #, Vec(X, n)) -> Vec(Y, n)
 mapVec(_, _, _, _, []) = []
-mapVec(_, _, f, succ(n), x ,- xs) = f(x) ,- mapVec(!, !, f, n, xs)
+mapVec(_, _, f, _, x ,- xs) = f(x) ,- mapVec(!, !, f, !, xs)

brat/examples/unified.brat

@@ -28,3 +28,15 @@ swapFront(X :: *, n :: #, Vec(X, n)) -> Vec(X, n)
 swapFront(_, _, []) = []
 swapFront(_, _, [x]) = [x]
 swapFront(X, _, cons(x, cons(y, zs))) = cons(y, cons(x, zs))
+
+filled(X :: *, n :: #, Vec(X, full(n))) -> Vec(X, full(n))
+filled(_, _, xsl =, x ,= xsr) = xsl =, x ,= xsr
+
+fullId(X :: *, n :: #, Vec(X, full(n))) -> Vec(X, full(n))
+fullId(_, _, [] =,= []) = []
+fullId(_, _, [] =, x ,= []) = [x]
+fullId(_, succ(n), xl =, x ,= xr) = fullId(!, n, xl) =, x ,= fullId(!, n, xr)
+
+-- mapAndConquer(X :: *, Y :: *, n :: #, f :: { X -> Y }, Vec(X, succ(n))) -> Vec(Y, succ(n))
+-- mapAndConquer(_, _, doub(n), f, xsl =, x ,= xsr) = mapAndConquer(!, !, n, f, xsl) =, f(x) ,= mapAndConquer(!, !, n, f, xsr)
+-- mapAndConquer(_, _, succ(doub(n)), f, xsl =,= xsr) = mapAndConquer(!, !, n, f, xsl) =,= mapAndConquer(!, !, n, f, xsr)