From 462e763175ca5ec3c4b1da0669bf703d2da1f137 Mon Sep 17 00:00:00 2001 From: lecopivo Date: Tue, 27 Aug 2024 04:00:02 +0200 Subject: [PATCH] derive handle for IndexType --- SciLean/Data/IndexType.lean | 158 +++++++++++++++++++++++++++++++----- 1 file changed, 138 insertions(+), 20 deletions(-) diff --git a/SciLean/Data/IndexType.lean b/SciLean/Data/IndexType.lean index 27a5ae28..f46847ad 100644 --- a/SciLean/Data/IndexType.lean +++ b/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