derive handle for IndexType

SciLean/Data/IndexType.lean

@@ -5,6 +5,8 @@ import Mathlib.Data.Fintype.Prod
 import Mathlib.Data.Fintype.Sum
 import Mathlib.Data.Fintype.Basic
 import Mathlib.Algebra.Order.GroupWithZero.Canonical
+import Mathlib.Tactic.ProxyType
+import Mathlib.Tactic.DeriveFintype
 
 import SciLean.Util.SorryProof
 import SciLean.Tactic.RefinedSimp
@@ -27,6 +29,30 @@ instance (a : List α) : Size a where
 
 inductive IndexType.Range (I : Type u)
   | full
+  -- endpoint
+  -- startpoint
+  -- interval
+
+@[inline]
+def IndexType.Range.ofEquiv (_ : I ≃ J) (i : Range I) : Range J :=
+  match i with
+  | .full => .full
+
+@[inline]
+def IndexType.Range.ofProd (i : Range (I×J)) : Range I × Range J :=
+  match i with
+  | .full => (.full, .full)
+
+@[inline]
+def IndexType.Range.prod (i : Range I) (j : Range J) : Range (I×J) :=
+  match i, j with
+  | .full, .full => .full
+
+@[inline]
+def IndexType.Range.ofSigma (i : Range ((_:I)×J)) : Range I × Range J :=
+  match i with
+  | .full => (.full, .full)
+
 
 inductive IndexType.Stream (I : Type u) where
   /-- Stream for range `range` that has not been started -/
@@ -36,6 +62,48 @@ inductive IndexType.Stream (I : Type u) where
   -- /-- Stream that has been exhausted -/
   -- | done
 
+@[inline]
+def IndexType.Stream.val! [Inhabited I] (s : Stream I) : I :=
+  match s with
+  | .start _ => panic! "can't take value of start stream!"
+  | .val i _ => i
+
+@[inline]
+def IndexType.Stream.ofEquiv (f : I ≃ J) (i : Stream I) : Stream J :=
+  match i with
+  | .start r => .start (r.ofEquiv f)
+  | .val i r => .val (f i) (r.ofEquiv f)
+
+@[inline]
+def IndexType.Stream.ofProd (i : Stream (I×J)) : Stream I × Stream J :=
+  match i with
+  | .start r =>
+    let (r,s) := r.ofProd
+    (.start r, .start s)
+  | .val (i,j) r =>
+    let (r, s) := r.ofProd
+    (.val i r, .val j s)
+
+@[inline]
+def IndexType.Stream.prod (i : Stream I) (j : Stream J) : Stream (I×J) :=
+  match i, j with
+  | .start ri, .start rj => .start (ri.prod rj)
+  | .start ri, .val _ rj => .start (ri.prod rj)
+  | .val _ ri, .start rj => .start (ri.prod rj)
+  | .val i ri, .val j rj => .val (i,j) (ri.prod rj)
+
+
+@[inline]
+def IndexType.Stream.ofSigma (i : Stream ((_:I)×J)) : Stream I × Stream J :=
+  match i with
+  | .start r =>
+    let (r,s) := r.ofSigma
+    (.start r, .start s)
+  | .val ⟨i,j⟩ r =>
+    let (r, s) := r.ofSigma
+    (.val i r, .val j s)
+
+
 -- This is alternative to LeanColls.IndexType which unfortunatelly has two universe parameters
 -- and because of that it is very difficult to work with.
 class IndexType (I : Type u) extends Fintype I, Stream (IndexType.Stream I) I, Size I where
@@ -50,7 +118,13 @@ def IndexType.first? (I : Type u) [IndexType I] : Option I :=
   else
     .none
 
-
+def IndexType.Range.first? {I : Type u} [IndexType I] (r : Range I) : Option I :=
+  match r with
+  | .full =>
+    if h : size I ≠ 0 then
+      .some (IndexType.fromFin ⟨0, by omega⟩)
+    else
+      .none
 
 
 instance : IndexType Empty where
@@ -107,25 +181,42 @@ instance {α β} [IndexType α] [IndexType β] : IndexType (α × β) where
     (IndexType.fromFin i, IndexType.fromFin j)
   next? s :=
     match s with
-    | .start _r =>
-      if let .some a := IndexType.first? α then
-        if let .some b := IndexType.first? β then
-          .some ((a,b), .val (a,b) .full) -- todo: split the range somehow
-        else
-          .none
-      else
-        .none
-    | .val (a,b) _r =>
-      if let .some sa := Stream.next? (IndexType.Stream.val a .full) then
-        .some ((sa.1,b), .val (sa.1,b) .full)
-      else
-        if let .some sb := Stream.next? (IndexType.Stream.val b .full) then
-          if let .some a := IndexType.first? α then
-            .some ((a,sb.1), .val (a,sb.1) .full)
-          else
-            .none
-        else
-          .none
+    | .start r =>
+      let (ri, rj) := r.ofProd
+      match ri.first?, rj.first? with
+      | .some i, .some j => .some ((i,j), .val (i,j) r)
+      | _, _ => .none
+    | .val (i,j) r =>
+      let (ri,rj) := r.ofProd
+      let si := IndexType.Stream.val i ri
+      let sj := IndexType.Stream.val j rj
+      match Stream.next? si with
+      | .some (i', si) => .some ((i',j), si.prod sj)
+      | .none =>
+        match ri.first?, Stream.next? sj with
+        | .some i', .some (j', sj) => .some ((i',j'), (IndexType.Stream.val i' ri).prod sj)
+        | _, _ => .none
+    -- match s with
+    -- | .start _r =>
+    --   if let .some a := IndexType.first? α then
+    --     if let .some b := IndexType.first? β then
+    --       .some (⟨a,b⟩, .val ⟨a,b⟩ .full) -- todo: split the range somehow
+    --     else
+    --       .none
+    --   else
+    --     .none
+    -- | .val ⟨a,b⟩ _r =>
+    --   if let .some sa := Stream.next? (IndexType.Stream.val a .full) then
+    --     .some (⟨sa.1,b⟩, .val ⟨sa.1,b⟩ .full)
+    --   else
+    --     if let .some sb := Stream.next? (IndexType.Stream.val b .full) then
+    --       if let .some a := IndexType.first? α then
+    --         .some (⟨a,sb.1⟩, .val ⟨a,sb.1⟩ .full)
+    --       else
+    --         .none
+    --     else
+    --       .none
+
 
 instance {α β} [IndexType α] [IndexType β] : IndexType ((_ : α) × β) where
   size := size α * size β
@@ -225,6 +316,33 @@ instance (a b : Int) : IndexType (Icc a b) where
         .none
 
 
+def IndexType.ofEquiv [IndexType I] [Fintype J] (f : I≃J) : IndexType J where
+  size := size I
+  toFin j := toFin (f.symm j)
+  fromFin idx := f (fromFin idx)
+  next? s :=
+    match Stream.next? (s.ofEquiv f.symm) with
+    | .some (i, si) => .some (f i, si.ofEquiv f)
+    | .none => .none
+
+
+open Lean Elab Command
+def mkIndexTypeInstanceHandler (declNames : Array Name) : CommandElabM Bool := do
+  if declNames.size != 1 then
+    return false -- mutually inductive types are not supported
+  let id : Ident := mkIdent declNames[0]!
+  try
+    elabCommand (← `(deriving instance Fintype for $id))
+  catch _ =>
+    pure ()
+
+  elabCommand (← `(instance : IndexType $id := IndexType.ofEquiv (proxy_equiv% $id)))
+  return true
+
+initialize
+  registerDerivingHandler ``IndexType mkIndexTypeInstanceHandler
+
+
 namespace SciLean
 end SciLean