From 8e90cc192f1f97f5734872b30fed89dcddee872b Mon Sep 17 00:00:00 2001
From: GiggleLiu <cacate0129@gmail.com>
Date: Tue, 2 Jul 2024 14:11:11 +0800
Subject: [PATCH] annotate a single line

---
 docs/src/man/registers.md        |  2 +-
 lib/YaoPlots/src/YaoPlots.jl     |  2 ++
 lib/YaoPlots/src/helperblock.jl  | 18 ++++++++++++++----
 lib/YaoPlots/src/vizcircuit.jl   | 12 ++++++++++++
 lib/YaoPlots/test/helperblock.jl |  7 +++++++
 5 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/docs/src/man/registers.md b/docs/src/man/registers.md
index 51c36f43..b0342119 100644
--- a/docs/src/man/registers.md
+++ b/docs/src/man/registers.md
@@ -140,7 +140,7 @@ We also have some "cheating" version [`measure`](@ref) that does not collapse st
 ```@repl register
 measure!(reg0, 1)  # measure the qubit, the state collapses
 measure!(reg0)  # measure all qubits
-measure(reg0, 3)  # measure the qubit 3 times, the state does not collapse (hacky)
+measure(reg0, 3)  # measure the qubit at location 3, the state does not collapse (hacky)
 reorder!(reg0, 7:-1:1)  # reorder the qubits
 measure!(reg0)
 invorder!(reg0)  # reverse the order of qubits
diff --git a/lib/YaoPlots/src/YaoPlots.jl b/lib/YaoPlots/src/YaoPlots.jl
index baa89090..931ee97a 100644
--- a/lib/YaoPlots/src/YaoPlots.jl
+++ b/lib/YaoPlots/src/YaoPlots.jl
@@ -8,10 +8,12 @@ import Thebes
 using Luxor: @layer, Point
 using Thebes: Point3D, project
 using LinearAlgebra: tr
+using YaoBlocks
 
 export CircuitStyles, vizcircuit, darktheme!, lighttheme!
 export bloch_sphere, BlochStyles
 export plot
+export LabelBlock, addlabel, LineAnnotation, line_annotation
 
 """An alias of `vizcircuit`"""
 plot(;kwargs...) = x->plot(x;kwargs...)
diff --git a/lib/YaoPlots/src/helperblock.jl b/lib/YaoPlots/src/helperblock.jl
index 057a92bb..84ccb4be 100644
--- a/lib/YaoPlots/src/helperblock.jl
+++ b/lib/YaoPlots/src/helperblock.jl
@@ -1,6 +1,3 @@
-using YaoBlocks
-export LabelBlock
-
 """
     LabelBlock{BT,D} <: TagBlock{BT,D}
 
@@ -31,7 +28,6 @@ Base.adjoint(x::LabelBlock) = LabelBlock(adjoint(content(x)), endswith(x.name, "
 Base.copy(x::LabelBlock) = LabelBlock(copy(content(x)), x.name, x.color)
 YaoBlocks.Optimise.to_basictypes(block::LabelBlock) = block
 
-export addlabel
 addlabel(b::AbstractBlock; name=string(b), color="transparent") = LabelBlock(b, name, color)
 
 # to fix issue 
@@ -49,3 +45,17 @@ function YaoBlocks.print_tree(
 )
    print(io, node.name)
 end
+
+"""
+    LineAnnotation{D} <: TrivialGate{D}
+"""
+struct LineAnnotation{D} <: TrivialGate{D}
+    name::String
+    color::String
+end
+line_annotation(name::String; color="black", nlevel=2) = LineAnnotation{nlevel}(name, color)
+
+Base.copy(x::LineAnnotation) = LineAnnotation(x.name, x.color)
+YaoBlocks.Optimise.to_basictypes(block::LineAnnotation) = block
+YaoBlocks.nqudits(::LineAnnotation) = 1
+YaoBlocks.print_block(io::IO, blk::LineAnnotation) = YaoBlocks.printstyled(io, blk.name; color=Symbol(blk.color))
\ No newline at end of file
diff --git a/lib/YaoPlots/src/vizcircuit.jl b/lib/YaoPlots/src/vizcircuit.jl
index 76348c0b..458d624c 100644
--- a/lib/YaoPlots/src/vizcircuit.jl
+++ b/lib/YaoPlots/src/vizcircuit.jl
@@ -374,6 +374,18 @@ function draw!(c::CircuitGrid, cb::LabelBlock, address, controls)
     CircuitStyles.gate_bgcolor[] = temp
 end
 
+function draw!(c::CircuitGrid, cb::LineAnnotation, address, controls)
+    @assert length(address) == 1 && isempty(controls) "LineAnnotation should be a single line, without control."
+    CircuitStyles.textcolor[], temp = cb.color, CircuitStyles.textcolor[]
+    _annotate!(c, address[1], cb.name)
+    CircuitStyles.textcolor[] = temp
+end
+function _annotate!(c::CircuitGrid, loc::Integer, name::AbstractString)
+    wspace, fontsize = text_width_and_size(name)
+    i = frontier(c, loc) + 0.1
+    CircuitStyles.render(CircuitStyles.Text(fontsize), (c[i, loc-0.2], name, wspace, fontsize))
+end
+
 # [:KronBlock, :RepeatedBlock, :CachedBlock, :Subroutine, :(YaoBlocks.AD.NoParams)]
 function draw!(c::CircuitGrid, p::CompositeBlock, address, controls)
     barrier_style = CircuitStyles.barrier_for_chain[]
diff --git a/lib/YaoPlots/test/helperblock.jl b/lib/YaoPlots/test/helperblock.jl
index a279634b..79a0ee17 100644
--- a/lib/YaoPlots/test/helperblock.jl
+++ b/lib/YaoPlots/test/helperblock.jl
@@ -27,3 +27,10 @@ using Test
     @test vizcircuit(c1) isa Drawing
     @test vizcircuit(c2) isa Drawing
 end
+
+@testset "annotate line" begin
+    ann = line_annotation("test"; color="red")
+    @test ann isa LineAnnotation
+    c = chain(put(5, 2=>X), put(5, 2=>ann),  put(5, 3=>Y))
+    @test vizcircuit(c) isa Drawing
+end
\ No newline at end of file