Correct typing of container extensions.

2 files changed, 19 insertions, 10 deletions

diff --git a/src/Obs/NormalForm/Normalise.idr b/src/Obs/NormalForm/Normalise.idr
index e36e57c..a95e725 100644
--- a/src/Obs/NormalForm/Normalise.idr
+++ b/src/Obs/NormalForm/Normalise.idr
@@ -251,8 +251,6 @@ doSem {container, predSort, pred, arg} = do
   let shapeVar = Nothing
   let positionVar = Nothing
 
-  let inputDomain = MkDecl inputVar container.input
-
   shape <- doShape {inputRel = relevance container.inputSort, arg}
   let shape = Sorted.weaken [_] shape
 
@@ -286,11 +284,10 @@ doSem {container, predSort, pred, arg} = do
   let f = add (LogNormalForm' ann) (Left $ pure next) (Right . There . There . There)
   codomain <- subst' pred f
 
-  pure $ Cnstr $ Pi
-    { domainSort = container.inputSort
-    , codomainSort = max container.shapeSort (container.positionSort ~> predSort)
-    , domain = inputDomain
-    , codomain = Cnstr $ Sigma
+  pure $ MkLambda
+    { var = inputVar
+    , domainRel = relevance container.inputSort
+    , body = Cnstr $ Sigma
       { indexSort = container.shapeSort
       , elementSort = container.positionSort ~> predSort
       , index = shapeIndex
@@ -903,3 +900,13 @@ containerNextType container shape position =
         }
       }
     }
+
+export
+containerSemType : (container : ContainerTy ctx) -> (predSort : Universe) -> TypeNormalForm ctx
+containerSemType container predSort =
+  Cnstr $ Pi
+    { domainSort = container.inputSort
+    , codomainSort = succ (max container.shapeSort (container.positionSort ~> predSort))
+    , domain = MkDecl Nothing container.input
+    , codomain = cast $ max container.shapeSort (container.positionSort ~> predSort)
+    }
diff --git a/src/Obs/Typing.idr b/src/Obs/Typing.idr
index e8db1db..50c4824 100644
--- a/src/Obs/Typing.idr
+++ b/src/Obs/Typing.idr
@@ -432,10 +432,12 @@ infer' ctx (Sem {pred = MkLambda var pred, arg}) = do
         }
   (predSort, pred) <- inferType predCtx pred
 
-  let semSort = container.inputSort ~> max container.shapeSort (container.positionSort ~> predSort)
-  semType <- doSem container predSort pred arg
+  let semType = containerSemType container predSort
+  sem <- doSem container predSort pred arg
 
-  pure (succ semSort ** (cast semSort, semType))
+  pure ((container.inputSort ~>
+         succ (max container.shapeSort (container.positionSort ~> predSort))) **
+        (semType, sem))
 
 infer' ctx Bool = do
     info "encountered bool"